diffusers/tools

1#!/usr/bin/env python32# !pip install transformers accelerate3import os4import PIL5from pathlib import Path6from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, DDIMScheduler, StableDiffusionInpaintPipeline, StableDiffusionImg2ImgPipeline, StableDiffusionControlNetImg2ImgPipeline7from diffusers.utils import load_image8import numpy as np9from huggingface_hub import HfApi10import torch11 12api = HfApi()13init_image = load_image(14    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy.png"15)16init_image = init_image.resize((512, 512))17 18generator = torch.Generator(device="cpu").manual_seed(33)19 20mask_image = load_image(21    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy_mask.png"22)23mask_image = mask_image.resize((512, 512))24 25 26def make_inpaint_condition(image, image_mask):27    image = np.array(image.convert("RGB")).astype(np.float32) / 255.028    image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.029 30    assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"31    image[image_mask > 0.5] = -1.0  # set as masked pixel32    image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)33    image = torch.from_numpy(image)34    return image35 36 37control_image = make_inpaint_condition(init_image, mask_image)38 39mask_image = PIL.Image.open("/home/patrick/images/mask.png").convert('RGB')40init_image = PIL.Image.open("/home/patrick/images/init.png").convert('RGB')41control_image = PIL.Image.open("/home/patrick/images/seg.png").convert('RGB')42 43controlnet = ControlNetModel.from_pretrained(44    "mfidabel/controlnet-segment-anything", torch_dtype=torch.float1645)46pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(47    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float1648)49 50# speed up diffusion process with faster scheduler and memory optimization51pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)52 53pipe.enable_model_cpu_offload()54 55# generate image56for t in [2]:57    image = pipe(58        "a bench in front of a beautiful lake and white mountain",59        num_inference_steps=t,60        generator=generator,61        eta=1.0,62        image=init_image,63        mask_image=mask_image,64        control_image=control_image,65    ).images[0]66 67    file_name = f"aa_{t}"68    path = os.path.join(Path.home(), "images", f"{file_name}.png")69    image.save(path)70 71    api.upload_file(72        path_or_fileobj=path,73        path_in_repo=path.split("/")[-1],74        repo_id="patrickvonplaten/images",75        repo_type="dataset",76    )77    print(f"https://huggingface.co/datasets/patrickvonplaten/images/blob/main/{file_name}.png")78