Scribd
Upload
11 views2 pages

Stable Diffusion Inpainting Guide

This document is a Python script that performs image inpainting using a trained Stable Diffusion model. It takes input parameters for model path, validation image, validation mask, output directory, and seed for reproducibility, and processes images to generate inpainted results. The script ensures that all images are resized to 256x256 pixels and applies preprocessing filters before saving the output images.

Uploaded by

areej
Copyright
© All Rights Reserved
We take content rights seriously. If you suspect this is your content, claim it here.
Available Formats
Download as TXT, PDF, TXT or read online on Scribd
11 views2 pages

Stable Diffusion Inpainting Guide

This document is a Python script that performs image inpainting using a trained Stable Diffusion model. It takes input parameters for model path, validation image, validation mask, output directory, and seed for reproducibility, and processes images to generate inpainted results. The script ensures that all images are resized to 256x256 pixels and applies preprocessing filters before saving the output images.

Uploaded by

areej
Copyright
© All Rights Reserved
We take content rights seriously. If you suspect this is your content, claim it here.
Available Formats
Download as TXT, PDF, TXT or read online on Scribd

import argparse

import os
import torch
from PIL import Image, ImageFilter
from diffusers import (
StableDiffusionInpaintPipeline,
UNet2DConditionModel,
DDPMScheduler
)
from transformers import CLIPTextModel

parser = [Link](description="Inference")
parser.add_argument("--model_path", type=str, required=True, help="Path to trained
model.")
parser.add_argument("--validation_image", type=str, required=True, help="Path to
validation image.")
parser.add_argument("--validation_mask", type=str, required=True, help="Path to
validation mask.")
parser.add_argument("--output_dir", type=str, default="./test-infer/", help="Output
directory.")
parser.add_argument("--seed", type=int, default=None, help="Seed for reproducible
inference.")

args = parser.parse_args()

if __name__ == "__main__":
[Link](args.output_dir, exist_ok=True)
generator = None

# Create & load model


pipe = StableDiffusionInpaintPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-inpainting",
torch_dtype=torch.float32,
revision=None
)

[Link] = UNet2DConditionModel.from_pretrained(args.model_path,
subfolder="unet", revision=None)
pipe.text_encoder = CLIPTextModel.from_pretrained(args.model_path,
subfolder="text_encoder", revision=None)
[Link] = DDPMScheduler.from_config([Link])

# Ensure model processes 256x256 images


[Link]["sample_size"] = 256

pipe = [Link]("cuda")

if [Link] is not None:


generator = [Link](device="cuda").manual_seed([Link])

# Ensure all images & masks are the same size (256x256)
target_size = (256, 256)
image = [Link](args.validation_image).convert("RGB").resize(target_size,
[Link])
mask_image = [Link](args.validation_mask).convert("L").resize(target_size,
[Link])

# Apply preprocessing filters


erode_kernel = [Link](3)
mask_image = mask_image.filter(erode_kernel)

blur_kernel = [Link](1)
mask_image = mask_image.filter(blur_kernel)

for idx in range(16):


result = pipe(
prompt="a photo of sks",
image=image,
mask_image=mask_image,
num_inference_steps=200,
guidance_scale=1,
generator=generator,
).images[0]

# Ensure result is also 256x256 before pasting


result = [Link](target_size, [Link])

# Composite final image


result = [Link](result, image, mask_image)
[Link](f"{args.output_dir}/{idx}.png")

# Clear GPU memory at the end


del pipe
[Link].empty_cache()

You might also like