add ddpm inversion

app.py

@@ -61,20 +61,22 @@ def prep(config):
         model_key = "stabilityai/stable-diffusion-2-depth"
     toy_scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler")
     toy_scheduler.set_timesteps(config["save_steps"])
-    print("config[save_steps]", config["save_steps"])
     timesteps_to_save, num_inference_steps = get_timesteps(toy_scheduler, num_inference_steps=config["save_steps"],
                                                            strength=1.0,
                                                            device=device)
-    print("YOOOO timesteps to save", timesteps_to_save)
+
 
     # seed_everything(config["seed"])
     if not config["frames"]: # original non demo setting
         save_path = os.path.join(config["save_dir"],
+                                 f'inversion_{config[inversion]}',
                                  f'sd_{config["sd_version"]}',
                                  Path(config["data_path"]).stem,
                                  f'steps_{config["steps"]}',
                                  f'nframes_{config["n_frames"]}') 
         os.makedirs(os.path.join(save_path, f'latents'), exist_ok=True)
+        if opt[inversion] == 'ddpm':
+            os.makedirs(os.path.join(save_path, f'latents'), exist_ok=True)
         add_dict_to_yaml_file(os.path.join(config["save_dir"], 'inversion_prompts.yaml'), Path(config["data_path"]).stem, config["inversion_prompt"])    
         # save inversion prompt in a txt file
         with open(os.path.join(save_path, 'inversion_prompt.txt'), 'w') as f:
@@ -82,43 +84,53 @@ def prep(config):
     else:
         save_path = None
     
-    model = Preprocess(device, config,
+    model = Preprocess(device, 
+                       config,
                       vae=vae,
                       text_encoder=text_encoder,
                       scheduler=scheduler,
                       tokenizer=tokenizer,
                       unet=unet)
-    print(type(model.config["batch_size"]))
-    frames, latents, total_inverted_latents, rgb_reconstruction = model.extract_latents(
+
+    frames_and_latents, rgb_reconstruction = model.extract_latents(
                                          num_steps=model.config["steps"],
                                          save_path=save_path,
                                          batch_size=model.config["batch_size"],
                                          timesteps_to_save=timesteps_to_save,
                                          inversion_prompt=model.config["inversion_prompt"],
+                                         inversion_type=model.config["inversion"],
+                                         skip_steps=model.config["skip_steps"],
+                                         reconstruction=model.config["reconstruct"]
     )
 
-    
-    return frames, latents, total_inverted_latents, rgb_reconstruction
+    if model.config["inversion"] == 'ddpm':
+        frames, latents, total_inverted_latents, zs = frames_and_latents
+        return frames, latents, total_inverted_latents, zs, rgb_reconstruction
+    else:
+        frames, latents, total_inverted_latents = frames_and_latents
+        return frames, latents, total_inverted_latents, rgb_reconstruction
+
     
 def preprocess_and_invert(input_video,
                           frames,
                           latents,
                           inverted_latents,
+                          zs,
                           seed, 
                           randomize_seed,
                           do_inversion,
-                          # save_dir: str = "latents",
                           steps,
                           n_timesteps = 50,
                           batch_size: int = 8,
                           n_frames: int = 40,
                           inversion_prompt:str = '',
+                          skip_steps: int = 15,
                           
               ):
     sd_version = "2.1"
-    height = 512
+    height: int = 512
     weidth: int = 512
-    print("n timesteps", n_timesteps)
+
     if do_inversion or randomize_seed:
         preprocess_config = {}
         preprocess_config['H'] = height
@@ -134,30 +146,37 @@ def preprocess_and_invert(input_video,
         preprocess_config['frames'] = video_to_frames(input_video)
         preprocess_config['data_path'] = input_video.split(".")[0]
         
+        preprocess_config['inversion'] = 'ddpm'
+        preprocess_config['skip_steps'] = skip_steps
+        preprocess_config['reconstruct'] = False
+        
+        
 
         if randomize_seed:
             seed = randomize_seed_fn()
         seed_everything(seed)
         
-        frames, latents, total_inverted_latents, rgb_reconstruction = prep(preprocess_config)
-        print(total_inverted_latents.keys())
-        print(len(total_inverted_latents.keys()))
-        frames = gr.State(value=frames)
-        latents = gr.State(value=latents)
-        inverted_latents = gr.State(value=total_inverted_latents)
+        frames, latents, total_inverted_latents, zs, rgb_reconstruction = prep(preprocess_config)
+        
+        frames = gr.State(value = frames)
+        latents = gr.State(value = latents)
+        inverted_latents = gr.State(value = total_inverted_latents)
+        zs = gr.State(value = zs)
         do_inversion = False
    
-    return frames, latents, inverted_latents, do_inversion
+    return frames, latents, inverted_latents, zs, do_inversion
 
 
 def edit_with_pnp(input_video,
                   frames, 
                   latents,
                   inverted_latents,
+                  zs,
                   seed,
                   randomize_seed,
                   do_inversion,
                   steps,
+                  skip_steps: int = 15,
                   prompt: str = "a marble sculpture of a woman running, Venus de Milo",
                   # negative_prompt: str = "ugly, blurry, low res, unrealistic, unaesthetic",
                   pnp_attn_t: float = 0.5,
@@ -183,14 +202,18 @@ def edit_with_pnp(input_video,
     config["pnp_attn_t"] = pnp_attn_t
     config["pnp_f_t"] = pnp_f_t
     config["pnp_inversion_prompt"] = inversion_prompt
+    config["inversion"] = "ddpm"
+    config["skip_steps"] = skip_steps
+    
     
     
     if do_inversion:
-        frames, latents, inverted_latents, do_inversion =  preprocess_and_invert(
+        frames, latents, inverted_latents, zs, do_inversion =  preprocess_and_invert(
                           input_video,
                           frames,
                           latents,
                           inverted_latents,
+                          zs,
                           seed, 
                           randomize_seed,
                           do_inversion,
@@ -198,7 +221,8 @@ def edit_with_pnp(input_video,
                           n_timesteps,
                           batch_size,
                           n_frames,
-                          inversion_prompt)
+                          inversion_prompt,
+                          skip_steps)
         do_inversion = False
         
     
@@ -207,12 +231,13 @@ def edit_with_pnp(input_video,
     seed_everything(seed)
     
     
-    editor = TokenFlow(config=config,pipe=tokenflow_pipe, frames=frames.value, inverted_latents=inverted_latents.value)
+    editor = TokenFlow(config=config,pipe=tokenflow_pipe, frames=frames.value, inverted_latents=inverted_latents.value, zs= zs.value)
     edited_frames = editor.edit_video()
-
-    save_video(edited_frames, 'tokenflow_PnP_fps_30.mp4', fps=n_fps)
+    
+    edit_video_path = f'tokenflow_PnP_fps_{n_fps}.mp4'
+    save_video(edited_frames, edit_video_path, fps=n_fps)
     # path = export_to_video(edited_frames)
-    return 'tokenflow_PnP_fps_30.mp4', frames, latents, inverted_latents, do_inversion
+    return edit_video_path, frames, latents, inverted_latents, zs, do_inversion
 
 ########
 # demo #
@@ -238,6 +263,7 @@ with gr.Blocks(css="style.css") as demo:
     frames = gr.State()
     inverted_latents = gr.State()
     latents = gr.State()
+    zs = gr.State()
     do_inversion = gr.State(value=True)
 
     with gr.Row():
@@ -252,15 +278,7 @@ with gr.Blocks(css="style.css") as demo:
                             label="Describe your edited video",
                             max_lines=1, value=""
                         )
-    # with gr.Group(visible=False) as share_btn_container:
-        # with gr.Group(elem_id="share-btn-container"):
-        #     community_icon = gr.HTML(community_icon_html, visible=True)
-        #     loading_icon = gr.HTML(loading_icon_html, visible=False)
-        #     share_button = gr.Button("Share to community", elem_id="share-btn", visible=True)
-        
-   
-    # with gr.Row():
-    #     inversion_progress = gr.Textbox(visible=False, label="Inversion progress")
+    
                
     with gr.Row():
         run_button = gr.Button("Edit your video!", visible=True)
@@ -274,8 +292,10 @@ with gr.Blocks(css="style.css") as demo:
                         randomize_seed = gr.Checkbox(label='Randomize seed', value=False)
                         gudiance_scale = gr.Slider(label='Guidance Scale', minimum=1, maximum=30,
                                               value=7.5, step=0.5, interactive=True)
-                        steps = gr.Slider(label='Inversion steps', minimum=10, maximum=500,
-                                              value=200, step=1, interactive=True)
+                        steps = gr.Slider(label='Inversion steps', minimum=10, maximum=200,
+                                              value=50, step=1, interactive=True)
+                        skip_steps = gr.Slider(label='Skip Steps', minimum=5, maximum=25,
+                                              value=5, step=1, interactive=True)
                         
                     with gr.Column(min_width=100):
                         inversion_prompt = gr.Textbox(lines=1, label="Inversion prompt", interactive=True, placeholder="")
@@ -284,7 +304,7 @@ with gr.Blocks(css="style.css") as demo:
                         n_frames = gr.Slider(label='Num frames', minimum=2, maximum=200,
                                               value=24, step=1, interactive=True)
                         n_timesteps = gr.Slider(label='Diffusion steps', minimum=25, maximum=100,
-                                              value=25, step=25, interactive=True)
+                                              value=50, step=25, interactive=True)
                         n_fps = gr.Slider(label='Frames per second', minimum=1, maximum=60,
                                               value=10, step=1, interactive=True)
                         
@@ -300,6 +320,11 @@ with gr.Blocks(css="style.css") as demo:
         fn = reset_do_inversion,
         outputs = [do_inversion],
         queue = False)
+    
+    steps.change(
+        fn = reset_do_inversion,
+        outputs = [do_inversion],
+        queue = False)
 
     inversion_prompt.change(
         fn = reset_do_inversion,
@@ -326,6 +351,7 @@ with gr.Blocks(css="style.css") as demo:
                       frames,
                       latents,
                       inverted_latents,
+                      zs,
                       seed, 
                       randomize_seed,
                       do_inversion,
@@ -333,11 +359,13 @@ with gr.Blocks(css="style.css") as demo:
                       n_timesteps,
                       batch_size,
                       n_frames,
-                      inversion_prompt
+                      inversion_prompt,
+                      skip_steps
           ],
           outputs = [frames,
                      latents,
                      inverted_latents,
+                     zs,
                      do_inversion
               
           ])
@@ -347,10 +375,12 @@ with gr.Blocks(css="style.css") as demo:
                                frames, 
                               latents,
                               inverted_latents,
+                              zs,
                               seed,
                               randomize_seed,
                               do_inversion,
                               steps,
+                              skip_steps,
                               prompt,                             
                               pnp_attn_t,
                               pnp_f_t,
@@ -360,7 +390,7 @@ with gr.Blocks(css="style.css") as demo:
                               gudiance_scale,
                               inversion_prompt,
                               n_fps ],
-                                 outputs = [output_video, frames, latents, inverted_latents, do_inversion]
+                                 outputs = [output_video, frames, latents, inverted_latents, zs, do_inversion]
                                 )
 
     gr.Examples(
@@ -371,4 +401,5 @@ with gr.Blocks(css="style.css") as demo:
     )
 
 demo.queue()
+
 demo.launch()