diffusers/tools

1#!/usr/bin/env python32import torch3import os4from huggingface_hub import HfApi5from pathlib import Path6from diffusers.utils import load_image7import cv28from PIL import Image9import numpy as np10 11from diffusers import (12    ControlNetModel,13    StableDiffusionControlNetImg2ImgPipeline,14    StableDiffusionControlNetInpaintPipeline,15    DiffusionPipeline,16    UniPCMultistepScheduler,17)18import sys19 20checkpoint = sys.argv[1]21 22# image = load_image(23#     "https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png"24# )25 26img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"27mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"28image = load_image(img_url).resize((512, 512))29mask_image = load_image(mask_url).resize((512, 512))30 31np_image = np.array(image)32 33low_threshold = 10034high_threshold = 20035 36np_image = cv2.Canny(np_image, low_threshold, high_threshold)37np_image = np_image[:, :, None]38np_image = np.concatenate([np_image, np_image, np_image], axis=2)39canny_image = Image.fromarray(np_image)40 41controlnet = ControlNetModel.from_pretrained(checkpoint, torch_dtype=torch.float16)42# pipe = DiffusionPipeline.from_pretrained(43#    "runwayml/stable-diffusion-inpainting", controlnet=controlnet, torch_dtype=torch.float16, custom_pipeline="stable_diffusion_controlnet_inpaint"44# )45pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(46    "runwayml/stable-diffusion-inpainting",47    controlnet=controlnet,48    torch_dtype=torch.float16,49)50 51pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)52pipe.enable_model_cpu_offload()53 54generator = torch.manual_seed(0)55text_prompt="a blue dog"56# out_image = pipe("A blue dog", num_inference_steps=50, generator=generator, image=image, mask_image=mask_image, controlnet_conditioning_image=canny_image).images[0]57out_image = pipe(58    text_prompt,59    num_inference_steps=20,60    generator=generator,61    image=image,62    mask_image=mask_image,63    control_image=canny_image,64).images[0]65 66path = os.path.join(Path.home(), "images", "aa.png")67out_image.save(path)68 69api = HfApi()70 71api.upload_file(72    path_or_fileobj=path,73    path_in_repo=path.split("/")[-1],74    repo_id="patrickvonplaten/images",75    repo_type="dataset",76)77print("https://huggingface.co/datasets/patrickvonplaten/images/blob/main/aa.png")78