Update app.py

app.py

@@ -16,6 +16,7 @@ from lvdm.utils.dist_utils import setup_dist, gather_data
 from lvdm.utils.saving_utils import npz_to_video_grid, npz_to_imgsheet_5d
 from utils import load_model, get_conditions, make_model_input_shape, torch_to_np
 from lvdm.models.modules.lora import change_lora
+from lvdm.utils.saving_utils import tensor_to_mp4
 
 from huggingface_hub import hf_hub_download
 
@@ -110,6 +111,50 @@ def sample_text2video(model, prompt, n_samples, batch_size,
     assert(all_videos.shape[0] >= n_samples)
     return all_videos
 
+def adapter_guided_synthesis(model, prompts, videos, noise_shape, sampler, n_samples=1, ddim_steps=50, ddim_eta=1., \
+                        unconditional_guidance_scale=1.0, unconditional_guidance_scale_temporal=None, **kwargs):
+    ddim_sampler = sampler 
+
+    batch_size = noise_shape[0]
+    ## get condition embeddings (support single prompt only)
+    if isinstance(prompts, str):
+        prompts = [prompts]
+    cond = model.get_learned_conditioning(prompts)
+    if unconditional_guidance_scale != 1.0:
+        prompts = batch_size * [""]
+        uc = model.get_learned_conditioning(prompts)
+    else:
+        uc = None
+    
+    ## adapter features: process in 2D manner
+    b, c, t, h, w = videos.shape
+    extra_cond = model.get_batch_depth(videos, (h,w))
+    features_adapter = model.get_adapter_features(extra_cond)
+
+    batch_variants = []
+    for _ in range(n_samples):
+        if ddim_sampler is not None:
+            samples, _ = ddim_sampler.sample(S=ddim_steps,
+                                            conditioning=cond,
+                                            batch_size=noise_shape[0],
+                                            shape=noise_shape[1:],
+                                            verbose=False,
+                                            unconditional_guidance_scale=unconditional_guidance_scale,
+                                            unconditional_conditioning=uc,
+                                            eta=ddim_eta,
+                                            temporal_length=noise_shape[2],
+                                            conditional_guidance_scale_temporal=unconditional_guidance_scale_temporal,
+                                            features_adapter=features_adapter,
+                                            **kwargs
+                                            )        
+        ## reconstruct from latent to pixel space
+        batch_images = model.decode_first_stage(samples, decode_bs=1, return_cpu=False)
+        batch_variants.append(batch_images)
+    ## variants, batch, c, t, h, w
+    batch_variants = torch.stack(batch_variants)
+    return batch_variants.permute(1, 0, 2, 3, 4, 5), extra_cond
+
+
 def save_results(videos, 
                  save_name="results", save_fps=8, save_mp4=True, 
                  save_npz=False, save_mp4_sheet=False, save_jpg=False
@@ -124,6 +169,9 @@ def save_results(videos,
         
     return os.path.join(save_subdir, f"{save_name}_{i:03d}.mp4")
 
+def save_results_control(batch_samples, batch_conds):
+    return 
+
 def get_video(prompt, seed, ddim_steps):
     seed_everything(seed)
     samples = sample_text2video(model, prompt, n_samples = 1, batch_size = 1,
@@ -156,6 +204,14 @@ def get_video_lora(prompt, seed, ddim_steps, model_choice):
                           )
     return save_results(samples)
 
+def get_video_control(prompt, input_video, seed, ddim_steps):
+    seed_everything(seed)
+    h,w = 512//8, 512//8
+    noise_shape = [1, model.channels, model.temporal_length,h,w]
+    batch_samples, batch_conds = adapter_guided_synthesis(model, prompt,input_video,noise_shape, sampler=ddim_sampler, n_samples = 1,
+                          ddim_steps=ddim_steps
+                          )
+    return save_results_control(batch_samples, batch_conds)
 
 from gradio_t2v import create_demo as create_demo_basic
 from gradio_videolora import create_demo as create_demo_videolora
@@ -170,6 +226,8 @@ with gr.Blocks(css='style.css') as demo:
             create_demo_basic(get_video)
         with gr.TabItem('VideoLoRA'):
             create_demo_videolora(get_video_lora)
+        with gr.TabItem('VideoControl'):
+            create_demo_videolora(get_video_control)
 
 demo.queue(api_open=False).launch()