Tight Inversion SDXL demo

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+**/__pycache__/

app.py

@@ -0,0 +1,230 @@
+import gradio as gr
+import torch
+import random
+import numpy as np
+from PIL import Image
+from diffusers.utils import load_image
+import yaml
+
+# Import parts of your original code
+from diffusers.utils.torch_utils import randn_tensor
+from transformers import CLIPVisionModelWithProjection
+
+# These functions are assumed to be available in your project:
+from src.config import RunConfig
+from src.enums import Model_Type, Scheduler_Type
+from src.utils.enums_utils import get_pipes, is_stochastic, model_type_to_size
+from src.utils.images_utils import crop_center_square_and_resize
+
+# Use the same weight name as in your code
+weight_name = "ip-adapter-plus_sdxl_vit-h.safetensors"
+
+# A helper to create noise (duplicated from your code)
+def create_noise_list(model_type, length, dtype, generator=None):
+    img_size = model_type_to_size(model_type)
+    VQAE_SCALE = 8
+    latents_size = (1, 4, img_size[0] // VQAE_SCALE, img_size[1] // VQAE_SCALE)
+    device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+    return [randn_tensor(latents_size, dtype=dtype, device=device, generator=generator) for i in range(length)]
+
+config_path = "configs/run_configs/ddim/ddim_ipa_scale_0.4.yaml"
+with open(config_path, 'r') as file:
+    cfg = RunConfig.from_yaml(yaml.safe_load(file))
+
+# Set device and seed
+device = "cuda" if torch.cuda.is_available() else "cpu"
+seed = cfg.seed if cfg.seed is not None else 42
+torch.manual_seed(seed)
+random.seed(seed)
+np.random.seed(seed)
+
+# Set torch dtype based on config
+dtype = torch.float32 if cfg.use_float32 else torch.bfloat16
+
+# Load image encoder if using IP-Adapter
+image_encoder = None
+if cfg.inference_use_ipa:
+    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
+        "h94/IP-Adapter",
+        subfolder="models/image_encoder",
+        torch_dtype=dtype,
+    ).to(device)
+
+# Get both inversion and inference (editing) pipelines.
+pipe_inversion, pipe_inference = get_pipes(cfg, image_encoder=image_encoder, device=device)
+# pipe_inversion.cfg = cfg
+# pipe_inference.cfg = cfg
+
+def edit_demo(source_image, source_prompt, edit_prompt, ipa_scale, guidance_scale, sharpening_factor, use_negative_prompt):
+    """
+    Given a source image, a source prompt and an edit prompt, this function:
+      1. Preprocesses the image.
+      2. Uses the inversion pipeline (with the source prompt) to compute image latents
+         using an IPA scale set by the user.
+      3. Uses the editing pipeline (with the edit prompt) to generate the edited image
+         using the user-provided guidance scale (while inversion guidance scale is fixed to 1)
+         and applies a sharpening factor.
+    """
+    if source_image is None:
+        return None
+
+    pipe_inversion, pipe_inference = get_pipes(cfg, image_encoder=image_encoder, device=device)
+    pipe_inversion.cfg = cfg
+    pipe_inference.cfg = cfg
+
+    # Preprocess the source image: convert to RGB and crop/resize
+    input_image = source_image.convert("RGB")
+    model_size = model_type_to_size(cfg.model_type)  # e.g. (512, 512)
+    if input_image.size == (500, 500):
+        # add noise to the image (magicbrush)
+        input_image = input_image.resize(model_size)
+        # # Convert image to numpy array
+        image_array = np.array(input_image)
+
+        # Define Gaussian noise parameters
+        mean = 0
+        stddev = 17.5  # Standard deviation (adjust for more/less noise)
+
+        # Generate Gaussian noise
+        gaussian_noise = np.random.normal(mean, stddev, image_array.shape).astype(np.float32)
+
+        # Add noise to the image
+        noisy_image_array = np.clip(image_array + gaussian_noise, 0, 255).astype(np.uint8)
+
+        # Convert back to PIL Image
+        noisy_image = Image.fromarray(noisy_image_array)
+        input_image = noisy_image
+    else:
+        input_image = crop_center_square_and_resize(input_image, model_size)
+
+    # Prepare IP adapter arguments using the user-specified IPA scale
+    editing_ipa_args = {}
+    inversion_ipa_args = {}
+    reconstruction_ipa_args = {}
+    if cfg.inference_use_ipa or cfg.inversion_use_ipa:
+        # pipe_inversion.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name=weight_name)
+        pipe_inference.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name=weight_name)
+        pipe_inference.set_ip_adapter_scale(cfg.inference_ipa_scale)
+        print("Preparing image embeds...")
+        image_embeds = pipe_inference.prepare_ip_adapter_image_embeds(
+            ip_adapter_image=input_image,
+            ip_adapter_image_embeds=None,
+            device=device,
+            num_images_per_prompt=1,
+            do_classifier_free_guidance=True,
+        )
+        if cfg.use_image_embeds_for_null_prompt:
+            image_embeds = [torch.stack([image_embeds[0][1], image_embeds[0][1]])]
+        inference_image_embeds = image_embeds
+        inversion_image_embeds = image_embeds
+        reconstruction_image_embeds = image_embeds
+        if cfg.guidance_scale == 1.0 or cfg.guidance_scale == 0.0:
+            inference_image_embeds = [image_embeds[0][None, 1]]
+        if cfg.inversion_guidance_scale == 1.0 or cfg.inversion_guidance_scale == 0.0:
+            inversion_image_embeds = [image_embeds[0][None, 1]]
+        if cfg.reconstruction_guidance_scale == 1.0 or cfg.reconstruction_guidance_scale == 0.0:
+            reconstruction_image_embeds = [image_embeds[0][None, 1]]
+        if cfg.inference_use_ipa:
+            editing_ipa_args = {"ip_adapter_image_embeds": inference_image_embeds}
+            reconstruction_ipa_args = {"ip_adapter_image_embeds": reconstruction_image_embeds}
+        if cfg.inversion_use_ipa:
+            inversion_ipa_args = {"ip_adapter_image_embeds": inversion_image_embeds}
+        pipe_inference.unload_ip_adapter()
+
+    # Create a generator with the seed
+    generator = torch.Generator(device=device).manual_seed(seed)
+    if is_stochastic(cfg.scheduler_type):
+        noise = create_noise_list(cfg.model_type, cfg.num_inversion_steps, dtype, generator=generator)
+        pipe_inversion.scheduler.set_noise_list(noise)
+        pipe_inference.scheduler.set_noise_list(noise)
+
+    # --- Inversion Stage ---
+    # Use the source prompt to invert the image into latents.
+    print("Performing inversion...")
+    if cfg.inversion_use_ipa:
+        pipe_inversion.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name=weight_name)
+        pipe_inversion.set_ip_adapter_scale(ipa_scale)  # user-specified IPA scale
+    inversion_result = pipe_inversion(
+        prompt=source_prompt if source_prompt != "" else None,
+        negative_prompt=cfg.negative_prompt,
+        num_inversion_steps=cfg.num_inversion_steps,
+        num_inference_steps=cfg.num_inversion_steps,
+        generator=generator,
+        image=input_image,
+        guidance_scale=1.0,  # inversion guidance scale is fixed to 1
+        use_cfgpp=cfg.use_cfgpp_inversion,
+        strength=cfg.inversion_max_step,
+        denoising_start=1.0 - cfg.inversion_max_step,
+        num_gd_steps=cfg.num_gd_steps,
+        gd_step_size=cfg.gd_step_size,
+        optimization_start=cfg.optimization_start,
+        normalize=cfg.normalize,
+        **inversion_ipa_args
+    )
+    # The inversion result returns a tuple; we take the first latent.
+    inv_latent = inversion_result[0][0]
+
+    # --- Editing Stage ---
+    print("Performing editing...")
+    negative_prompt = cfg.negative_prompt if not use_negative_prompt else source_prompt
+    if cfg.inference_use_ipa:
+        pipe_inference.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name=weight_name)
+        pipe_inference.set_ip_adapter_scale(ipa_scale)  # user-specified IPA scale
+    edited_image = pipe_inference(
+        image=inv_latent,
+        prompt=edit_prompt if edit_prompt != "" else None,
+        negative_prompt=negative_prompt if negative_prompt != "" else None,
+        denoising_start=1.0 - cfg.inversion_max_step,
+        strength=cfg.inversion_max_step,
+        num_inference_steps=cfg.num_inference_steps,
+        guidance_scale=guidance_scale,  # use the user-specified editing guidance scale
+        guidance_rescale=cfg.guidance_rescale,
+        use_cfgpp=cfg.use_cfgpp_inference,
+        sharpening_factor=sharpening_factor,
+        **editing_ipa_args
+    ).images[0]
+
+    pipe_inference.unload_ip_adapter()
+    pipe_inversion.unload_ip_adapter()
+
+    return edited_image
+
+# Build Gradio Interface with additional inputs for IPA scale, guidance scale, and sharpening factor.
+demo = gr.Interface(
+    fn=edit_demo,
+    inputs=[
+        gr.Image(type="pil", label="Source Image"),
+        gr.Textbox(lines=2, placeholder="Enter source prompt here", label="Source Prompt"),
+        gr.Textbox(lines=2, placeholder="Enter edit prompt here", label="Edit Prompt"),
+        gr.Slider(minimum=0.1, maximum=1.0, step=0.05, value=0.4, label="IPA Scale"),
+        gr.Slider(minimum=1.0, maximum=20.0, step=0.1, value=7.5, label="Guidance Scale for Editing"),
+        gr.Slider(minimum=1.0, maximum=3.0, step=0.1, value=1.5, label="Sharpening Factor"),
+        gr.Checkbox(label="Use negative prompt", value=False),
+    ],
+    outputs=gr.Image(type="pil", label="Edited Image"),
+    title="Tight Inversion SDXL Demo",
+    description=(
+        "Upload an image, provide a source prompt (for inversion) and an edit prompt, "
+        "set the IPA scale (for both inversion and editing), the guidance scale for editing "
+        "(inversion guidance scale is fixed to 1), and the guidance scale sharpening factor (editing only), "
+        "then view the edited image. You can start with the provided examples."
+    ),
+    examples=[
+        # Example template:
+        # ["path/to/example_image.jpg", "A sunny landscape", "A stormy night", 0.4, 7.5, 1.5],
+        # Add more examples here.
+        ["example_images/animals/7.jpg", "A photo of a husky", "A photo of a cat", 0.4, 7.5, 1.0, False],
+        ["example_images/tests/diner_square.png", "a photo of people dining in a diner", "a photo of robots dining in a diner", 0.4, 7.5, 1.0, True],
+        ["example_images/tests/lion.jpg", "a lion in the field", "a lion made of lego in the field", 0.4, 7.5, 1.0, True],
+        ["example_images/garibis_images/5.jpg", "metal elephant statues", "real living elephants", 0.3, 4.0, 1.0, True],
+        # ["example_images/tests/monkey.jpg", "a photo of a monkey sitting on a branch in the forest", "a photo of a beaver sitting on a branch in the forest", 0.7, 12, 1.5, False],
+        ["example_images/editAR_images/dog_forest.jpg", "a dog running in the forest", "a forest with no one around", 0.4, 7.5, 1.0, False],
+        ["example_images/animals/0.jpg", "A photo of a gazelle", "A photo of a gazelle wearing a red hat", 0.4, 7.5, 1.5, False],
+        # ["example_images/garibis_images/1.jpg", "white christmas tree", "christmas tree with red ornaments", 0.55, 7.5, 2.0, False],
+        ["example_images/people/5.jpg", "a person", "a person with a large thick beard", 0.5, 7.5, 2.0, True],
+        ["example_images/magic_brush/262283-input.png", "a lot of vases filled with lots of flowers", "a lot of vases filled with red roses", 0.4, 12, 1.4, True],
+    ]
+)
+
+if __name__ == "__main__":
+    demo.launch()

configs/run_configs/ddim/ddim_ipa_scale_0.4.yaml

@@ -0,0 +1,23 @@
+method: ddim
+model_type: SDXL
+scheduler_type: DDIM
+seed: 7865
+negative_prompt: null
+num_inference_steps: 50
+num_inversion_steps: 50
+inversion_max_step: 1.0
+guidance_scale: 7.5
+inversion_guidance_scale: 1.0
+reconstruction_guidance_scale: 1.0
+random_image_guidance_scale: 5.0
+inference_use_ipa: true
+inversion_use_ipa: true
+inference_ipa_scale: 0.4
+inversion_ipa_scale: 0.4
+saturation_removal_ipa_scale: 0.4
+num_gd_steps: 0
+gd_step_size: 0.001
+optimization_start: 0
+normalize: false
+random_inference_times: 0
+remove_cfg_saturation: false