test app_ddim

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 
-from app_haircolor_inpaint import create_demo as create_demo_haircolor
+from app_ddim import create_demo as create_demo_haircolor
 
 with gr.Blocks(css="style.css") as demo:
     with gr.Tabs():

app_ddim.py

@@ -0,0 +1,260 @@
+import spaces
+import gradio as gr
+import time
+import torch
+import numpy as np
+
+from tqdm.auto import tqdm
+from torchvision import transforms as tfms
+from PIL import Image
+from segment_utils import(
+    segment_image,
+    restore_result,
+)
+from diffusers import (
+    StableDiffusionPipeline,
+    DDIMScheduler,
+)
+
+BASE_MODEL = "stable-diffusion-v1-5/stable-diffusion-v1-5"
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+DEFAULT_INPUT_PROMPT = "a woman"
+DEFAULT_EDIT_PROMPT = "a woman with linen-blonde-hair"
+
+DEFAULT_CATEGORY = "hair"
+
+basepipeline = StableDiffusionPipeline.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+)
+
+basepipeline.scheduler = DDIMScheduler.from_config(basepipeline.scheduler.config)
+
+basepipeline = basepipeline.to(DEVICE)
+
+basepipeline.enable_model_cpu_offload()
+
+@spaces.GPU(duration=30)
+def image_to_image(
+    input_image: Image,
+    input_image_prompt: str,
+    edit_prompt: str,
+    num_steps: int,
+    start_step: int,
+    guidance_scale: float,
+):
+    run_task_time = 0
+    time_cost_str = ''
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+
+    with torch.no_grad():
+        latent = basepipeline.vae.encode(tfms.functional.to_tensor(input_image).unsqueeze(0).to(DEVICE) * 2 - 1)
+    l = 0.18215 * latent.latent_dist.sample()
+    inverted_latents = invert(l, input_image_prompt, num_inference_steps=num_steps)
+    generated_image = sample(
+        edit_prompt,
+        start_latents=inverted_latents[-(start_step + 1)][None],
+        start_step=start_step,
+        num_inference_steps=num_steps,
+        guidance_scale=guidance_scale,
+    )[0]
+    
+    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+
+    return generated_image, time_cost_str
+
+def make_inpaint_condition(image, image_mask):
+    image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
+    image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0
+
+    assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
+    image[image_mask > 0.5] = -1.0  # set as masked pixel
+    image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)
+    image = torch.from_numpy(image)
+    return image
+
+## Inversion
+@torch.no_grad()
+def invert(
+    start_latents,
+    prompt,
+    guidance_scale=3.5,
+    num_inference_steps=80,
+    num_images_per_prompt=1,
+    do_classifier_free_guidance=True,
+    negative_prompt="",
+    device=DEVICE,
+):
+
+    # Encode prompt
+    text_embeddings = basepipeline._encode_prompt(
+        prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt
+    )
+
+    # Latents are now the specified start latents
+    latents = start_latents.clone()
+
+    # We'll keep a list of the inverted latents as the process goes on
+    intermediate_latents = []
+
+    # Set num inference steps
+    basepipeline.scheduler.set_timesteps(num_inference_steps, device=device)
+
+    # Reversed timesteps <<<<<<<<<<<<<<<<<<<<
+    timesteps = reversed(basepipeline.scheduler.timesteps)
+
+    for i in tqdm(range(1, num_inference_steps), total=num_inference_steps - 1):
+
+        # We'll skip the final iteration
+        if i >= num_inference_steps - 1:
+            continue
+
+        t = timesteps[i]
+
+        # Expand the latents if we are doing classifier free guidance
+        latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+        latent_model_input = basepipeline.scheduler.scale_model_input(latent_model_input, t)
+
+        # Predict the noise residual
+        noise_pred = basepipeline.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
+
+        # Perform guidance
+        if do_classifier_free_guidance:
+            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+        current_t = max(0, t.item() - (1000 // num_inference_steps))  # t
+        next_t = t  # min(999, t.item() + (1000//num_inference_steps)) # t+1
+        alpha_t = basepipeline.scheduler.alphas_cumprod[current_t]
+        alpha_t_next = basepipeline.scheduler.alphas_cumprod[next_t]
+
+        # Inverted update step (re-arranging the update step to get x(t) (new latents) as a function of x(t-1) (current latents)
+        latents = (latents - (1 - alpha_t).sqrt() * noise_pred) * (alpha_t_next.sqrt() / alpha_t.sqrt()) + (
+            1 - alpha_t_next
+        ).sqrt() * noise_pred
+
+        # Store
+        intermediate_latents.append(latents)
+
+    return torch.cat(intermediate_latents)
+
+# Sample function (regular DDIM)
+@torch.no_grad()
+def sample(
+    prompt,
+    start_step=0,
+    start_latents=None,
+    guidance_scale=3.5,
+    num_inference_steps=30,
+    num_images_per_prompt=1,
+    do_classifier_free_guidance=True,
+    negative_prompt="",
+    device=DEVICE,
+):
+
+    # Encode prompt
+    text_embeddings = basepipeline._encode_prompt(
+        prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt
+    )
+
+    # Set num inference steps
+    basepipeline.scheduler.set_timesteps(num_inference_steps, device=device)
+
+    # Create a random starting point if we don't have one already
+    if start_latents is None:
+        start_latents = torch.randn(1, 4, 64, 64, device=device)
+        start_latents *= basepipeline.scheduler.init_noise_sigma
+
+    latents = start_latents.clone()
+
+    for i in tqdm(range(start_step, num_inference_steps)):
+
+        t = basepipeline.scheduler.timesteps[i]
+
+        # Expand the latents if we are doing classifier free guidance
+        latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+        latent_model_input = basepipeline.scheduler.scale_model_input(latent_model_input, t)
+
+        # Predict the noise residual
+        noise_pred = basepipeline.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
+
+        # Perform guidance
+        if do_classifier_free_guidance:
+            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+        # Normally we'd rely on the scheduler to handle the update step:
+        # latents = pipe.scheduler.step(noise_pred, t, latents).prev_sample
+
+        # Instead, let's do it ourselves:
+        prev_t = max(1, t.item() - (1000 // num_inference_steps))  # t-1
+        alpha_t = basepipeline.scheduler.alphas_cumprod[t.item()]
+        alpha_t_prev = basepipeline.scheduler.alphas_cumprod[prev_t]
+        predicted_x0 = (latents - (1 - alpha_t).sqrt() * noise_pred) / alpha_t.sqrt()
+        direction_pointing_to_xt = (1 - alpha_t_prev).sqrt() * noise_pred
+        latents = alpha_t_prev.sqrt() * predicted_x0 + direction_pointing_to_xt
+
+    # Post-processing
+    images = basepipeline.decode_latents(latents)
+    images = basepipeline.numpy_to_pil(images)
+
+    return images
+
+def get_time_cost(run_task_time, time_cost_str):
+    now_time = int(time.time()*1000)
+    if run_task_time == 0:
+        time_cost_str = 'start'
+    else:
+        if time_cost_str != '': 
+            time_cost_str += f'-->'
+        time_cost_str += f'{now_time - run_task_time}'
+    run_task_time = now_time
+    return run_task_time, time_cost_str
+
+def create_demo() -> gr.Blocks:
+    with gr.Blocks() as demo:
+        croper = gr.State()
+        with gr.Row():
+            with gr.Column():
+                input_image_prompt = gr.Textbox(lines=1, label="Input Image Prompt", value=DEFAULT_INPUT_PROMPT)
+                edit_prompt = gr.Textbox(lines=1, label="Edit Prompt", value=DEFAULT_EDIT_PROMPT)
+            with gr.Column():
+                num_steps = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Num Steps")
+                start_step = gr.Slider(minimum=0, maximum=100, value=15, step=1, label="Start Step")
+                guidance_scale = gr.Slider(minimum=0, maximum=30, value=5, step=0.5, label="Guidance Scale")
+            with gr.Column():
+                generate_size = gr.Number(label="Generate Size", value=512)
+                with gr.Accordion("Advanced Options", open=False):
+                    mask_expansion = gr.Number(label="Mask Expansion", value=50, visible=True)
+                    mask_dilation = gr.Slider(minimum=0, maximum=10, value=2, step=1, label="Mask Dilation")
+                    category = gr.Textbox(label="Category", value=DEFAULT_CATEGORY, visible=False)
+                g_btn = gr.Button("Edit Image")
+                
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(label="Input Image", type="pil")
+            with gr.Column():
+                restored_image = gr.Image(label="Restored Image", type="pil", interactive=False)
+            with gr.Column():
+                origin_area_image = gr.Image(label="Origin Area Image", type="pil", interactive=False)
+                generated_image = gr.Image(label="Generated Image", type="pil", interactive=False)
+                generated_cost = gr.Textbox(label="Time cost by step (ms):", visible=True, interactive=False)
+        
+        g_btn.click(
+            fn=segment_image,
+            inputs=[input_image, category, generate_size, mask_expansion, mask_dilation],
+            outputs=[origin_area_image, croper],
+        ).success(
+            fn=image_to_image,
+            inputs=[origin_area_image, input_image_prompt, edit_prompt, num_steps, start_step, guidance_scale],
+            outputs=[generated_image, generated_cost],
+        ).success(
+            fn=restore_result,
+            inputs=[croper, category, generated_image],
+            outputs=[restored_image],
+        )
+
+    return demo

requirements.txt

@@ -8,4 +8,5 @@ mediapipe
 spaces
 sentencepiece
 controlnet_aux
-peft
+peft
+tqdm