WompUniversity/Inpaint-Anything-no-errors

1import os2import sys3import glob4import argparse5import torch6import numpy as np7import PIL.Image as Image8from pathlib import Path9from diffusers import StableDiffusionInpaintPipeline10from utils.mask_processing import crop_for_filling_pre, crop_for_filling_post11from utils.crop_for_replacing import recover_size, resize_and_pad12from utils import load_img_to_array, save_array_to_img13 14 15def fill_img_with_sd(16        img: np.ndarray,17        mask: np.ndarray,18        text_prompt: str,19        device="cuda"20):21    pipe = StableDiffusionInpaintPipeline.from_pretrained(22        "stabilityai/stable-diffusion-2-inpainting",23        torch_dtype=torch.float32,24    ).to(device)25    img_crop, mask_crop = crop_for_filling_pre(img, mask)26    img_crop_filled = pipe(27        prompt=text_prompt,28        image=Image.fromarray(img_crop),29        mask_image=Image.fromarray(mask_crop)30    ).images[0]31    img_filled = crop_for_filling_post(img, mask, np.array(img_crop_filled))32    return img_filled33 34 35def replace_img_with_sd(36        img: np.ndarray,37        mask: np.ndarray,38        text_prompt: str,39        step: int = 50,40        device="cuda"41):42    pipe = StableDiffusionInpaintPipeline.from_pretrained(43        "stabilityai/stable-diffusion-2-inpainting",44        torch_dtype=torch.float32,45    ).to(device)46    img_padded, mask_padded, padding_factors = resize_and_pad(img, mask)47    img_padded = pipe(48        prompt=text_prompt,49        image=Image.fromarray(img_padded),50        mask_image=Image.fromarray(255 - mask_padded),51        num_inference_steps=step,52    ).images[0]53    height, width, _ = img.shape54    img_resized, mask_resized = recover_size(55        np.array(img_padded), mask_padded, (height, width), padding_factors)56    mask_resized = np.expand_dims(mask_resized, -1) / 25557    img_resized = img_resized * (1-mask_resized) + img * mask_resized58    return img_resized59 60 61def setup_args(parser):62    parser.add_argument(63        "--input_img", type=str, required=True,64        help="Path to a single input img",65    )66    parser.add_argument(67        "--text_prompt", type=str, required=True,68        help="Text prompt",69    )70    parser.add_argument(71        "--input_mask_glob", type=str, required=True,72        help="Glob to input masks",73    )74    parser.add_argument(75        "--output_dir", type=str, required=True,76        help="Output path to the directory with results.",77    )78    parser.add_argument(79        "--seed", type=int,80        help="Specify seed for reproducibility.",81    )82    parser.add_argument(83        "--deterministic", action="store_true",84        help="Use deterministic algorithms for reproducibility.",85    )86 87if __name__ == "__main__":88    """Example usage:89    python lama_inpaint.py \90        --input_img FA_demo/FA1_dog.png \91        --input_mask_glob "results/FA1_dog/mask*.png" \92        --text_prompt "a teddy bear on a bench" \93        --output_dir results 94    """95    parser = argparse.ArgumentParser()96    setup_args(parser)97    args = parser.parse_args(sys.argv[1:])98    device = "cuda" if torch.cuda.is_available() else "cpu"99 100    if args.deterministic:101        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"102        torch.use_deterministic_algorithms(True)103 104    img_stem = Path(args.input_img).stem105    mask_ps = sorted(glob.glob(args.input_mask_glob))106    out_dir = Path(args.output_dir) / img_stem107    out_dir.mkdir(parents=True, exist_ok=True)108 109    img = load_img_to_array(args.input_img)110    for mask_p in mask_ps:111        if args.seed is not None:112            torch.manual_seed(args.seed)113        mask = load_img_to_array(mask_p)114        img_filled_p = out_dir / f"filled_with_{Path(mask_p).name}"115        img_filled = fill_img_with_sd(116            img, mask, args.text_prompt, device=device)117        save_array_to_img(img_filled, img_filled_p)