diffusers/tools

1#!/usr/bin/env python32from huggingface_hub import HfApi3from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, StableDiffusionInpaintPipeline4from pathlib import Path5import os6import sys7import requests8from PIL import Image9from io import BytesIO10 11org_name = "lykon-models"12org_name = "lykon-absolute-realism"13 14file_path = sys.argv[1]15file_name = file_path.split("/")[-1].split(".safetensors")[0].replace("_", "-").replace(".", "-").lower()16 17if file_name.endswith("-"):18    file_name = file_name[:-1]19 20def download_image(url):21    response = requests.get(url)22    return Image.open(BytesIO(response.content)).convert("RGB")23 24kwargs = {}25if "xl" in file_name or "alpha2" in file_name:26    pipe = StableDiffusionXLPipeline.from_single_file(file_path)27elif "inpaint" in file_name:28    pipe = StableDiffusionInpaintPipeline.from_single_file(file_path)29    img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"30    mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"31 32    init_image = download_image(img_url).resize((512, 512))33    mask_image = download_image(mask_url).resize((512, 512))34    kwargs = {"image": init_image, "mask_image": mask_image}35else:36    pipe = StableDiffusionPipeline.from_single_file(file_path)37 38print("Test...")39prompt = "A lion in galaxies, spirals, nebulae, stars, smoke, iridescent, intricate detail, octane render, 8k"40pipe.to("cuda")41images = pipe(prompt=prompt, num_inference_steps=25, **kwargs).images42pipe.to("cpu")43 44api = HfApi()45for i, image in enumerate(images):46    image_file_name = f"bb_1_{i}"47    path = os.path.join(Path.home(), "images", f"{image_file_name}.png")48    image.save(path)49 50    api.upload_file(51        path_or_fileobj=path,52        path_in_repo=path.split("/")[-1],53        repo_id="patrickvonplaten/images",54        repo_type="dataset",55    )56    print(f"https://huggingface.co/datasets/patrickvonplaten/images/blob/main/{image_file_name}.png")57 58print("Upload...")59model_id = os.path.join(org_name, file_name)60pipe.push_to_hub(model_id, private=True)61