Update app.py

app.py

@@ -1,5 +1,10 @@
+import gradio as gr
 import torch
 from diffusers import StableDiffusionPipeline, DDIMScheduler
+from utils import *
+
+
+
 
 
 
@@ -9,45 +14,59 @@ model_id = "stabilityai/stable-diffusion-2-1-base"
 inv_pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to(device)
 inv_pipe.scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")
 
-
-def preprocess(data_path:str = 'examples/woman-running.mp4',
-              height:int = 512, 
-              weidth: int = 512,
-              # save_dir: str = "latents",
-              steps: int = 500,
-               batch_size: int = 8,
-               save_steps: int = 50,
-               n_frames: int = 40,
-               inversion_prompt:str = ''
+def randomize_seed_fn():
+    seed = random.randint(0, np.iinfo(np.int32).max)
+    return seed
+
+def preprocess_and_invert(video,
+                          frames,
+                          latents,
+                          inverted_latents,
+                          seed, 
+                          randomize_seed,
+                          do_inversion,
+                          height:int = 512, 
+                          weidth: int = 512,
+                          # save_dir: str = "latents",
+                          steps: int = 500,
+                          batch_size: int = 8,
+                          # save_steps: int = 50,
+                          n_frames: int = 40,
+                          inversion_prompt:str = ''
               ):
+
+    if do_inversion or randomize_seed:
         
-    # save_video_frames(data_path, img_size=(height, weidth))
-    frames = video_to_frames(data_path, img_size=(height, weidth))
-    # data_path = os.path.join('data', Path(video_path).stem)
-
-    toy_scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")
-    toy_scheduler.set_timesteps(save_steps)
-    timesteps_to_save, num_inference_steps = get_timesteps(toy_scheduler, num_inference_steps=save_steps,
-                                                           strength=1.0,
-                                                           device=device)
-    seed_everything(1)
+        # save_video_frames(data_path, img_size=(height, weidth))
+        frames = video_to_frames(video, img_size=(height, weidth))
+        # data_path = os.path.join('data', Path(video_path).stem)
     
-    frames, latents = get_data(inv_pipe, frames, n_frames)
-    # inverted_latents = noisy_latents
-    inverted_latents = extract_latents(inv_pipe, num_steps = steps, 
-                                       latent_frames = latents, 
-                                       batch_size = batch_size, 
-                                       timesteps_to_save = timesteps_to_save,
-                                       inversion_prompt = inversion_prompt,)
+        toy_scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")
+        toy_scheduler.set_timesteps(save_steps)
+        timesteps_to_save, num_inference_steps = get_timesteps(toy_scheduler, num_inference_steps=save_steps,
+                                                               strength=1.0,
+                                                               device=device)
+        if randomize_seed:
+            seed = randomize_seed_fn()
+        seed_everything(seed)
+        
+        frames, latents = get_data(inv_pipe, frames, n_frames)
+
+        inverted_latents = extract_latents(inv_pipe, num_steps = steps, 
+                                           latent_frames = latents, 
+                                           batch_size = batch_size, 
+                                           timesteps_to_save = timesteps_to_save,
+                                           inversion_prompt = inversion_prompt,)
+        frames = gr.State(value=frames)
+        latents = gr.State(value=latents)
+        inverted_latents = gr.State(value=inverted_latents)
+        do_inversion = False
    
 
+    return frames, latents, inverted_latents, do_inversion
 
 
 
-    return frames, latents, inverted_latents
-
-import gradio as gr
-
 ########
 # demo #
 ########
@@ -64,7 +83,10 @@ intro = """
 with gr.Blocks(css="style.css") as demo:
     
     gr.HTML(intro)
-
+    frames = gr.State()
+    inverted_latents = gr.State()
+    latents = gr.State()
+    do_inversion = gr.State(value=True)
 
     with gr.Row():
         input_vid = gr.Video(label="Input Video", interactive=True, elem_id="input_video")
@@ -79,16 +101,57 @@ with gr.Blocks(css="style.css") as demo:
         #     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():
+    #     inversion_progress = gr.Textbox(visible=False, label="Inversion progress")
+               
     with gr.Row():
         run_button = gr.Button("Edit your video!", visible=True)
-        
 
-    # with gr.Accordion("Advanced Options", open=False):
+    with gr.Accordion("Advanced Options", open=False):
+      with gr.Tabs() as tabs:
+
+          with gr.TabItem('General options', id=2):
+            with gr.Row():
+                with gr.Column(min_width=100):
+                    seed = gr.Number(value=0, precision=0, label="Seed", interactive=True)
+                    randomize_seed = gr.Checkbox(label='Randomize seed', value=False)
+                    steps = gr.Slider(label='Inversion steps', minimum=100, maximum=500,
+                                          value=500, step=1, interactive=True)
+                with gr.Column(min_width=100):
+                    inversion_prompt = gr.Textbox(lines=1, label="Inversion prompt", interactive=True, placeholder="")
+                    batch_size = gr.Slider(label='Batch size', minimum=1, maximum=10,
+                                          value=8, step=1, interactive=True)
+                    n_frames = gr.Slider(label='Num frames', minimum=20, maximum=200,
+                                          value=40, step=1, interactive=True)
+                    
+    input_vid.change(
+        fn = reset_do_inversion,
+        outputs = [do_inversion],
+        queue = False)
+
+    input_vid.upload(
+        fn = reset_do_inversion,
+        outputs = [do_inversion],
+        queue = False)
+    ).then(fn = preprocess_and_invert,
+          inputs = [input_vid,
+                      frames,
+                      latents,
+                      inverted_latents,
+                      seed, 
+                      randomize_seed,
+                      do_inversion,
+                      steps,
+                      batch_size,
+                      n_frames,
+                      inversion_prompt
+          ],
+          outputs = [frames,
+                     latents,
+                     inverted_latents,
+                     do_inversion
+              
+          ])