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" | |