Create app.py

app.py

@@ -0,0 +1,109 @@
+import gradio as gr
+import os
+from torch.optim import AdamW
+from diffusers import StableDiffusionPipeline
+from torch import autocast, inference_mode
+import torch
+import numpy as np
+
+from scheduling_ddim import DDIMScheduler
+
+
+device = 'cuda'
+# don't forget to add your token or comment if already logged in
+pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", 
+                                               scheduler=DDIMScheduler(beta_end=0.012,
+                                                                       beta_schedule="scaled_linear",
+                                                                       beta_start=0.00085),
+                                               use_auth_token="").to(device)
+_ = pipe.vae.requires_grad_(False)
+_ = pipe.text_encoder.requires_grad_(False)
+_ = pipe.unet.requires_grad_(False)
+                                               
+def preprocess(image):
+    w, h = image.size
+    w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
+    image = image.resize((w, h), resample=Image.LANCZOS)
+    image = np.array(image).astype(np.float32) / 255.0
+    image = image[None].transpose(0, 3, 1, 2)
+    image = torch.from_numpy(image)
+    return 2.0 * image - 1.0
+
+def im2latent(pipe, im, generator):
+    init_image = preprocess(im).to(pipe.device)
+    init_latent_dist = pipe.vae.encode(init_image).latent_dist
+    init_latents = init_latent_dist.sample(generator=generator)
+    
+    return init_latents * 0.18215
+
+
+def image_mod(init_image, source_prompt, prompt, scale, steps, seed):
+    # fix seed
+    g = torch.Generator(device=pipe.device).manual_seed(84)
+    
+    image_latents = im2latent(pipe, init_image, g)
+    pipe.scheduler.set_timesteps(steps)
+    # use text describing an image
+    # source_prompt = "a photo of a woman"
+    context = pipe._encode_prompt(source_prompt, pipe.device, 1, False, "")
+    
+    decoded_latents = image_latents.clone()
+    with autocast(device), inference_mode():
+        # we are pivoting timesteps as we are moving in opposite direction
+        timesteps = pipe.scheduler.timesteps.flip(0)
+        # this would be our targets for pivoting
+        init_trajectory = torch.empty(len(timesteps), *decoded_latents.size()[1:], device=decoded_latents.device, dtype=decoded_latents.dtype)
+        for i, t in enumerate(tqdm(timesteps)):
+            init_trajectory[i:i+1] = decoded_latents
+            noise_pred = pipe.unet(decoded_latents, t, encoder_hidden_states=context).sample
+            decoded_latents = pipe.scheduler.reverse_step(noise_pred, t, decoded_latents).next_sample
+            
+     # we would need to flip trajectory values for pivoting in right direction
+    init_trajectory = init_trajectory.cpu().flip(0)
+    
+    latents = decoded_latents.clone()
+    context_uncond = pipe._encode_prompt("", pipe.device, 1, False, "")
+    # we will be optimizing uncond text embedding
+    context_uncond.requires_grad_(True)
+    
+    # use same text
+    # prompt = "a photo of a woman"
+    context_cond = pipe._encode_prompt(prompt, pipe.device, 1, False, "")
+    
+    # default lr works
+    opt = AdamW([context_uncond])
+    
+    # concat latents for classifier-free guidance
+    latents = torch.cat([latents, latents])
+    latents.requires_grad_(True)
+    context = torch.cat((context_uncond, context_cond))
+    
+    with autocast(device):
+        for i, t in enumerate(tqdm(pipe.scheduler.timesteps)):
+            latents = pipe.scheduler.scale_model_input(latents, t)
+            uncond, cond = pipe.unet(latents, t, encoder_hidden_states=context).sample.chunk(2)
+            with torch.enable_grad():
+                latents = pipe.scheduler.step(uncond + scale * (cond - uncond), t, latents, generator=g).prev_sample
+            
+            opt.zero_grad()
+            # optimize uncond text emb
+            pivot_value = init_trajectory[[i]].to(pipe.device)
+            (latents - pivot_value).mean().backward()
+            opt.step()
+            latents = latents.detach()
+    
+    images = pipe.decode_latents(latents)
+    im = pipe.numpy_to_pil(images)[0]
+    return im
+
+
+demo = gr.Interface(
+    image_mod, 
+    inputs=[gr.Image(type="pil"), gr.Textbox("a photo of a person"), gr.Textbox("a photo of a person"), gr.Slider(0, 10, 0.5, 0.1), gr.Slider(0, 100, 51, 1), gr.Number(42)], 
+    outputs="image",
+    flagging_options=["blurry", "incorrect", "other"], examples=[
+        os.path.join(os.path.dirname(__file__), "images/00001.jpg"),
+        ])
+
+if __name__ == "__main__":
+    demo.launch()