Spaces:
Build error
Build error
File size: 1,141 Bytes
c92867b |
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 |
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"
|