Spaces:
Build error
Build error
File size: 1,537 Bytes
9ae63a0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
from diffusers import DiffusionPipeline, DDIMScheduler
from PIL import Image
import imageio
import torch
# https://huggingface.co/spaces/Manjushri/SD-2.0-Inpainting-CPU/blob/main/app.py
def resize(height,img):
baseheight = height
img = Image.open(img)
hpercent = (baseheight/float(img.size[1]))
wsize = int((float(img.size[0])*float(hpercent)))
img = img.resize((wsize,baseheight), Image.Resampling.LANCZOS)
return img
def img_preprocces(source_img, prompt, negative_prompt):
imageio.imwrite("data.png", source_img["image"])
imageio.imwrite("data_mask.png", source_img["mask"])
src = resize(512, "data.png")
src.save("src.png")
mask = resize(512, "data_mask.png")
mask.save("mask.png")
return src, mask
def stable_diffusion_inpaint(
image_path:str,
model_path:str,
prompt:str,
negative_prompt:str,
guidance_scale:int,
num_inference_step:int,
):
image, mask_image = img_preprocces(image_path, prompt, negative_prompt)
pipe = DiffusionPipeline.from_pretrained(
model_path,
revision="fp16",
torch_dtype=torch.float16,
)
pipe.to('cuda')
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
pipe.enable_xformers_memory_efficient_attention()
output = pipe(
prompt = prompt,
image = image,
mask_image=mask_image,
negative_prompt = negative_prompt,
num_inference_steps = num_inference_step,
guidance_scale = guidance_scale,
).images
return output
|