# Minimal inpainting inference script for SD1.5 and SDXL inpainting models. # インペインティングモデル（SD1.5およびSDXL）を使った最低限の推論スクリプト。 # # Loads a 9-channel inpainting UNet checkpoint, encodes a source image and mask # via VAE, then runs a standard denoising loop with the 9-channel UNet input: # [noisy_latents(4ch), mask(1ch), masked_image_latents(4ch)] # # Usage / 使い方: # # SD1.5 inpainting # python inpainting_minimal_inference.py \ # --ckpt_path sd-v1-5-inpainting.ckpt \ # --image input.png \ # --mask mask.png \ # --prompt "a yawning cat" # # # SDXL inpainting # python inpainting_minimal_inference.py \ # --ckpt_path sd_xl_base_1.0_inpainting.safetensors \ # --sdxl \ # --image input.png \ # --mask mask.png \ # --prompt "a yawning cat" # # # Generate a random procedural mask instead of supplying one # python inpainting_minimal_inference.py \ # --ckpt_path sd-v1-5-inpainting.ckpt \ # --image input.png \ # --prompt "a yawning cat" import argparse import datetime import math import os import random import numpy as np import torch from diffusers import EulerDiscreteScheduler from PIL import Image from tqdm import tqdm from transformers import CLIPTokenizer from library.device_utils import get_preferred_device, init_ipex init_ipex() from library import model_util, sdxl_model_util from library.utils import setup_logging setup_logging() import logging logger = logging.getLogger(__name__) # ------------------------------------------------------------------------- # Constants # ------------------------------------------------------------------------- SCHEDULER_LINEAR_START = 0.00085 SCHEDULER_LINEAR_END = 0.012 SCHEDULER_TIMESTEPS = 1000 SCHEDULER_SCHEDULE = "scaled_linear" SD15_VAE_SCALE_FACTOR = 0.18215 SDXL_VAE_SCALE_FACTOR = sdxl_model_util.VAE_SCALE_FACTOR # 0.13025 # ------------------------------------------------------------------------- # Helpers # ------------------------------------------------------------------------- def _timestep_embedding(timesteps, dim, max_period=10000): half = dim // 2 freqs = torch.exp( -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half ).to(device=timesteps.device) args = timesteps[:, None].float() * freqs[None] return torch.cat([torch.cos(args), torch.sin(args)], dim=-1) def _get_timestep_embedding(x, outdim): b = x.shape[0] x = torch.flatten(x) emb = _timestep_embedding(x, outdim) return torch.reshape(emb, (b, -1)) def load_image(path: str, width: int, height: int) -> torch.Tensor: """Load an RGB image and return a [-1, 1] float32 tensor of shape (1, 3, H, W).""" img = Image.open(path).convert("RGB").resize((width, height), Image.LANCZOS) arr = np.array(img).astype(np.float32) / 127.5 - 1.0 return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0) def load_mask(path: str, width: int, height: int) -> torch.Tensor: """ Load a mask image and return a binary float32 tensor of shape (1, 1, H, W). White pixels (>= 128) indicate the region to regenerate. """ mask = Image.open(path).convert("L").resize((width, height), Image.NEAREST) arr = np.array(mask).astype(np.float32) / 255.0 arr = (arr >= 0.5).astype(np.float32) return torch.from_numpy(arr).unsqueeze(0).unsqueeze(0) def make_default_mask(width: int, height: int, seed: int = None) -> torch.Tensor: """ Generate a wobbly-ellipse mask — a single connected organic region that is well-suited for inpainting sampling and inference. (The full random_mask used during training produces fragmented cloud patterns that can confuse the sampler when large proportions of the image are masked.) """ from library.mask_generator import wobbly_ellipse_mask pil_mask = wobbly_ellipse_mask(width, height, seed=seed) arr = (np.array(pil_mask).astype(np.float32) / 255.0 >= 0.5).astype(np.float32) return torch.from_numpy(arr).unsqueeze(0).unsqueeze(0) def encode_image(vae, image_t: torch.Tensor, vae_scale_factor: float, device, dtype): """Encode an image tensor to latents.""" image_t = image_t.to(device=device, dtype=dtype) with torch.no_grad(): latents = vae.encode(image_t).latent_dist.sample() return latents * vae_scale_factor def downsample_mask(mask_t: torch.Tensor, latent_h: int, latent_w: int, device, dtype): """Downsample a (1,1,H,W) mask to latent spatial size.""" mask_latent = torch.nn.functional.interpolate( mask_t, size=(latent_h, latent_w), mode="nearest" ) return mask_latent.to(device=device, dtype=dtype) # ------------------------------------------------------------------------- # SD1.5 text encoding # ------------------------------------------------------------------------- def encode_text_sd15(tokenizer, text_encoder, text: str, device, dtype): tokens = tokenizer( text, truncation=True, padding="max_length", max_length=tokenizer.model_max_length, return_tensors="pt", )["input_ids"].to(device) with torch.no_grad(): return text_encoder(tokens)[0].to(dtype) # ------------------------------------------------------------------------- # SDXL text encoding (mirrors sdxl_minimal_inference.py) # ------------------------------------------------------------------------- def encode_text_sdxl(tokenizer1, text_model1, tokenizer2, text_model2, text: str, device, dtype, original_size, crop_top_left, target_size): # encoder 1 tokens1 = tokenizer1( text, truncation=True, padding="max_length", max_length=tokenizer1.model_max_length, return_tensors="pt", )["input_ids"].to(device) with torch.no_grad(): enc1 = text_model1(tokens1, output_hidden_states=True, return_dict=True) emb1 = enc1["hidden_states"][11].to(dtype) # encoder 2 tokens2 = tokenizer2( text, truncation=True, padding="max_length", max_length=tokenizer2.model_max_length, return_tensors="pt", )["input_ids"].to(device) with torch.no_grad(): enc2 = text_model2(tokens2, output_hidden_states=True, return_dict=True) emb2 = enc2["hidden_states"][-2].to(dtype) pool2 = enc2["text_embeds"].to(dtype) text_emb = torch.cat([emb1, emb2], dim=2) # vector embedding oh, ow = original_size ct, cl = crop_top_left th, tw = target_size e1 = _get_timestep_embedding(torch.FloatTensor([[oh, ow]]), 256) e2 = _get_timestep_embedding(torch.FloatTensor([[ct, cl]]), 256) e3 = _get_timestep_embedding(torch.FloatTensor([[th, tw]]), 256) vec = torch.cat([pool2, e1.to(device, dtype), e2.to(device, dtype), e3.to(device, dtype)], dim=1) return text_emb, vec # ------------------------------------------------------------------------- # Main # ------------------------------------------------------------------------- if __name__ == "__main__": parser = argparse.ArgumentParser(description="Minimal inpainting inference for SD1.5 and SDXL inpainting models") parser.add_argument("--ckpt_path", type=str, required=True, help="Path to inpainting model checkpoint (.ckpt or .safetensors)") parser.add_argument("--image", type=str, required=True, help="Source image path (PNG/JPG)") parser.add_argument("--mask", type=str, default=None, help="Mask image path — white = region to regenerate. Omit to use a random procedural mask.") parser.add_argument("--prompt", type=str, default="", help="Positive prompt") parser.add_argument("--negative_prompt", type=str, default="", help="Negative prompt") parser.add_argument("--width", type=int, default=512, help="Output width (must be multiple of 64)") parser.add_argument("--height", type=int, default=512, help="Output height (must be multiple of 64)") parser.add_argument("--steps", type=int, default=30, help="Number of denoising steps") parser.add_argument("--guidance_scale", type=float, default=7.5, help="Classifier-free guidance scale") parser.add_argument("--seed", type=int, default=None, help="Random seed") parser.add_argument("--sdxl", action="store_true", help="Use SDXL model (default: SD1.5)") parser.add_argument("--output_dir", type=str, default=".", help="Directory to save output images") parser.add_argument( "--lora_weights", type=str, nargs="*", default=[], help="LoRA weights to merge, each as 'path;multiplier'", ) args = parser.parse_args() # Round dimensions to multiples of 64 args.width = max(64, args.width - args.width % 64) args.height = max(64, args.height - args.height % 64) DEVICE = get_preferred_device() DTYPE = torch.bfloat16 os.makedirs(args.output_dir, exist_ok=True) # ------------------------------------------------------------------ # Load model # ------------------------------------------------------------------ if args.sdxl: logger.info("Loading SDXL inpainting model...") text_model1, text_model2, vae, unet, _, _ = sdxl_model_util.load_models_from_sdxl_checkpoint( sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, args.ckpt_path, "cpu" ) VAE_SCALE = SDXL_VAE_SCALE_FACTOR tokenizer1 = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") tokenizer2 = CLIPTokenizer.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k") text_model1.to(DEVICE, dtype=DTYPE).eval() text_model2.to(DEVICE, dtype=DTYPE).eval() else: logger.info("Loading SD1.5 inpainting model...") text_encoder, vae, unet = model_util.load_models_from_stable_diffusion_checkpoint( False, args.ckpt_path, "cpu" ) VAE_SCALE = SD15_VAE_SCALE_FACTOR tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") text_encoder.to(DEVICE, dtype=DTYPE).eval() # fp16 VAE is unstable; use float32 vae_dtype = torch.float32 if DTYPE == torch.float16 else DTYPE vae.to(DEVICE, dtype=vae_dtype).eval() unet.to(DEVICE, dtype=DTYPE).eval() unet.set_use_memory_efficient_attention(True, False) # LoRA import networks.lora as lora for weights_file in args.lora_weights: if ";" in weights_file: weights_file, multiplier = weights_file.split(";") multiplier = float(multiplier) else: multiplier = 1.0 text_encoders = [text_model1, text_model2] if args.sdxl else [text_encoder] lora_model, weights_sd = lora.create_network_from_weights( multiplier, weights_file, vae, text_encoders, unet, None, True ) lora_model.merge_to(text_encoders, unet, weights_sd, DTYPE, DEVICE) # ------------------------------------------------------------------ # Scheduler # ------------------------------------------------------------------ scheduler = EulerDiscreteScheduler( num_train_timesteps=SCHEDULER_TIMESTEPS, beta_start=SCHEDULER_LINEAR_START, beta_end=SCHEDULER_LINEAR_END, beta_schedule=SCHEDULER_SCHEDULE, ) # ------------------------------------------------------------------ # Prepare image and mask # ------------------------------------------------------------------ if args.seed is not None: random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) image_t = load_image(args.image, args.width, args.height) if args.mask is not None: mask_t = load_mask(args.mask, args.width, args.height) logger.info(f"Using mask: {args.mask}") else: mask_t = make_default_mask(args.width, args.height, seed=args.seed) logger.info("Using wobbly-ellipse mask") # Encode masked image: source × (1 − mask) masked_image_t = image_t * (1.0 - mask_t) latent_h = args.height // 8 latent_w = args.width // 8 image_latents = encode_image(vae, image_t, VAE_SCALE, DEVICE, vae_dtype) masked_latents = encode_image(vae, masked_image_t, VAE_SCALE, DEVICE, vae_dtype) mask_latent = downsample_mask(mask_t, latent_h, latent_w, DEVICE, DTYPE) # Initial noise latents = torch.randn( (1, 4, latent_h, latent_w), device="cpu", dtype=torch.float32 ).to(DEVICE, dtype=DTYPE) * scheduler.init_noise_sigma # ------------------------------------------------------------------ # Text encoding # ------------------------------------------------------------------ if args.sdxl: size_args = dict( original_size=(args.height, args.width), crop_top_left=(0, 0), target_size=(args.height, args.width), ) c_emb, c_vec = encode_text_sdxl(tokenizer1, text_model1, tokenizer2, text_model2, args.prompt, DEVICE, DTYPE, **size_args) uc_emb, uc_vec = encode_text_sdxl(tokenizer1, text_model1, tokenizer2, text_model2, args.negative_prompt, DEVICE, DTYPE, **size_args) cond_args = (c_emb, c_vec) uncond_args = (uc_emb, uc_vec) def unet_call(inp, t, emb, vec): return unet(inp, t, emb, vec) else: c_emb = encode_text_sd15(tokenizer, text_encoder, args.prompt, DEVICE, DTYPE) uc_emb = encode_text_sd15(tokenizer, text_encoder, args.negative_prompt, DEVICE, DTYPE) cond_args = (c_emb,) uncond_args = (uc_emb,) def unet_call(inp, t, emb): return unet(inp, t, emb).sample # ------------------------------------------------------------------ # Denoising loop # ------------------------------------------------------------------ logger.info(f"Denoising ({args.steps} steps, guidance={args.guidance_scale})...") scheduler.set_timesteps(args.steps, DEVICE) timesteps = scheduler.timesteps.to(DEVICE) # Pre-generate fixed noise used to re-noise the original image latents each # step so that unmasked regions stay anchored to the original image. # After each scheduler step we composite: # latents = mask * denoised + (1 - mask) * (original + noise * sigma_next) # At the final step sigma_next = 0, so the unmasked area is exactly the original. compositing_noise = torch.randn_like(image_latents).to(DTYPE) image_latents_dtype = image_latents.to(DTYPE) mask_latent_dtype = mask_latent.to(DTYPE) masked_latents_dtype = masked_latents.to(DTYPE) with torch.no_grad(): for i, t in enumerate(tqdm(timesteps)): latent_model_input = scheduler.scale_model_input(latents, t) unet_input = torch.cat([latent_model_input, mask_latent_dtype, masked_latents_dtype], dim=1) if args.sdxl: noise_pred_uncond = unet(unet_input, t, uc_emb, uc_vec) noise_pred_cond = unet(unet_input, t, c_emb, c_vec) else: noise_pred_uncond = unet(unet_input, t, uc_emb).sample noise_pred_cond = unet(unet_input, t, c_emb).sample noise_pred = noise_pred_uncond + args.guidance_scale * (noise_pred_cond - noise_pred_uncond) latents = scheduler.step(noise_pred, t, latents).prev_sample # Composite: keep unmasked regions anchored to the original. # sigma_next is the noise level after this step (0.0 on the last step). sigma_next = scheduler.sigmas[i + 1].to(device=DEVICE, dtype=DTYPE) noised_original = image_latents_dtype + compositing_noise * sigma_next latents = mask_latent_dtype * latents + (1.0 - mask_latent_dtype) * noised_original # Decode latents = latents.to(vae_dtype) / VAE_SCALE image_out = vae.decode(latents).sample image_out = (image_out / 2 + 0.5).clamp(0, 1) image_np = image_out.cpu().permute(0, 2, 3, 1).float().numpy() image_np = (image_np * 255).round().astype("uint8") pil_images = [Image.fromarray(im) for im in image_np] timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") for i, img in enumerate(pil_images): out_path = os.path.join(args.output_dir, f"inpaint_{timestamp}_{i:03d}.png") img.save(out_path) logger.info(f"Saved: {out_path}") logger.info("Done!")