Spaces:
Running
Running
| import imageio | |
| import numpy as np | |
| from PIL import Image | |
| from diffusers import AutoPipelineForInpainting | |
| import torch | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"Using device for I2I: {device}") | |
| # Load the inpainting pipeline | |
| pipe = AutoPipelineForInpainting.from_pretrained( | |
| "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", | |
| torch_dtype=torch.float16, variant="fp16").to(device) | |
| def resize_image(image, height, width): | |
| """Resize image tensor to the desired height and width.""" | |
| return torch.nn.functional.interpolate(image, size=(height, width), mode='nearest') | |
| def dummy(img): | |
| """Save the composite image and generate a mask from the alpha channel.""" | |
| imageio.imwrite("output_image.png", img["composite"]) | |
| # Extract alpha channel from the first layer to create the mask | |
| alpha_channel = img["layers"][0][:, :, 3] | |
| mask = np.where(alpha_channel == 0, 0, 255).astype(np.uint8) | |
| return img["background"], mask | |
| def I2I(prompt, image, width=1024, height=1024, guidance_scale=8.0, num_inference_steps=20, strength=0.99): | |
| img_url, mask = dummy(image) | |
| # Resize image and mask to the target dimensions (height x width) | |
| img_url = Image.fromarray(img_url, mode="RGB").resize((width, height)) | |
| mask_url = Image.fromarray(mask,mode="L").resize((width, height)) | |
| # Make sure both image and mask are converted into correct tensors | |
| generator = torch.Generator(device=device).manual_seed(0) | |
| # Generate the inpainted image | |
| output = pipe( | |
| prompt=prompt, | |
| image=img_url, | |
| mask_image=mask_url, | |
| guidance_scale=guidance_scale, | |
| num_inference_steps=num_inference_steps, # steps between 15 and 30 work well for us | |
| strength=strength, # make sure to use `strength` below 1.0 | |
| generator=generator, | |
| ) | |
| return output.images[0] | |