Spaces:
Build error
Build error
import os | |
import torch | |
from PIL import Image | |
from diffusers import StableDiffusionInpaintPipeline | |
from diffusers import AutoencoderKL | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
pipe = StableDiffusionInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting", | |
torch_dtype=torch.float16, | |
revision="fp16", | |
vae=AutoencoderKL.from_pretrained( | |
"stabilityai/sd-vae-ft-mse", | |
torch_dtype=torch.float16 | |
).to(device) | |
).to(device) | |
pipe.enable_xformers_memory_efficient_attention() | |
os.makedirs("inpainting_output", exist_ok=True) | |
def inpaint(inputs, prompt): | |
image = inputs["image"].resize((image.size[0] - image.size[0] % 64, image.size[1] - image.size[1] % 64), Image.ANTIALIAS) | |
mask = inputs["mask"].resize((mask.size[0] - mask.size[0] % 64, mask.size[1] - mask.size[1] % 64), Image.ANTIALIAS) | |
output = pipe(prompt=prompt, image=image, mask_image=mask, guidance_scale=7.5, height=image.size[1], width=image.size[0]) | |
output.images[0].save(f"inpainting_output/output.png") | |
return output.images[0], "inpainting_output/output.png" | |