Update app.py

app.py

@@ -8,6 +8,7 @@ from pyramid_dit import PyramidDiTForVideoGeneration
 from diffusers.utils import export_to_video
 
 import spaces 
+import uuid
 
 import subprocess
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
@@ -62,49 +63,62 @@ model = load_model()
 
 # Text-to-video generation function
 @spaces.GPU(duration=240)
-def generate_video(prompt, duration, guidance_scale, video_guidance_scale):
+def generate_video(image, prompt, duration, guidance_scale, video_guidance_scale):
     temp = int(duration * 2.4)  # Convert seconds to temp value (assuming 24 FPS)
     torch_dtype = torch.bfloat16 if MODEL_DTYPE == "bf16" else torch.float32
-    
-    with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
-        frames = model.generate(
-            prompt=prompt,
-            num_inference_steps=[20, 20, 20],
-            video_num_inference_steps=[10, 10, 10],
-            height=768,
-            width=1280,
-            temp=temp,
-            guidance_scale=guidance_scale,
-            video_guidance_scale=video_guidance_scale,
-            output_type="pil",
-            save_memory=True,
-        )
-    
-    output_path = "output_video.mp4"
+    if(image):
+        cropped_image = center_crop(image, 1280, 720)
+        resized_image = cropped_image.resize((1280, 720))
+        with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
+            frames = model.generate_i2v(
+                prompt=prompt,
+                input_image=resized_image,
+                num_inference_steps=[10, 10, 10],
+                temp=temp,
+                guidance_scale=7.0,
+                video_guidance_scale=video_guidance_scale,
+                output_type="pil",
+                save_memory=True,
+            )
+    else:
+        with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
+            frames = model.generate(
+                prompt=prompt,
+                num_inference_steps=[20, 20, 20],
+                video_num_inference_steps=[10, 10, 10],
+                height=768,
+                width=1280,
+                temp=temp,
+                guidance_scale=guidance_scale,
+                video_guidance_scale=video_guidance_scale,
+                output_type="pil",
+                save_memory=True,
+            )
+    output_path = f"{str(uuid.uuid4())}_output_video.mp4"
     export_to_video(frames, output_path, fps=24)
     return output_path
 
 # Image-to-video generation function
-@spaces.GPU(duration=240)
-def generate_video_from_image(image, prompt, duration, video_guidance_scale):
-    temp = int(duration * 2.4)  # Convert seconds to temp value (assuming 24 FPS)
-    torch_dtype = torch.bfloat16 if MODEL_DTYPE == "bf16" else torch.float32
-    
-    target_size = (1280, 720)
-    cropped_image = center_crop(image, 1280, 720)
-    resized_image = cropped_image.resize((1280, 720))
-    
-    with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
-        frames = model.generate_i2v(
-            prompt=prompt,
-            input_image=resized_image,
-            num_inference_steps=[10, 10, 10],
-            temp=temp,
-            guidance_scale=7.0,
-            video_guidance_scale=video_guidance_scale,
-            output_type="pil",
-            save_memory=True,
-        )
+#@spaces.GPU(duration=240)
+#def generate_video_from_image(image, prompt, duration, video_guidance_scale):
+#    temp = int(duration * 2.4)  # Convert seconds to temp value (assuming 24 FPS)
+#    torch_dtype = torch.bfloat16 if MODEL_DTYPE == "bf16" else torch.float32
+#    
+#    target_size = (1280, 720)
+#    cropped_image = center_crop(image, 1280, 720)
+#    resized_image = cropped_image.resize((1280, 720))
+#    
+#    with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
+#        frames = model.generate_i2v(
+#            prompt=prompt,
+#            input_image=resized_image,
+#            num_inference_steps=[10, 10, 10],
+#            temp=temp,
+#            guidance_scale=7.0,
+#            video_guidance_scale=video_guidance_scale,
+#            output_type="pil",
+#            save_memory=True,
+#        )
     
     output_path = "output_video_i2v.mp4"
     export_to_video(frames, output_path, fps=24)
@@ -114,38 +128,41 @@ def generate_video_from_image(image, prompt, duration, video_guidance_scale):
 with gr.Blocks() as demo:
     gr.Markdown("# Pyramid Flow Video Generation Demo")
     
-    with gr.Tab("Text-to-Video"):
-        with gr.Row():
-            with gr.Column():
-                t2v_prompt = gr.Textbox(label="Prompt")
+    #with gr.Tab("Text-to-Video"):
+    with gr.Row():
+        with gr.Column():
+            with gr.Accordion("Image to Video (optional)", open=False):
+                i2v_image = gr.Image(type="pil", label="Input Image")
+            t2v_prompt = gr.Textbox(label="Prompt")
+            with gr.Accordion("Advanced settings", open=False):
                 t2v_duration = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Duration (seconds)")
                 t2v_guidance_scale = gr.Slider(minimum=1, maximum=15, value=9, step=0.1, label="Guidance Scale")
                 t2v_video_guidance_scale = gr.Slider(minimum=1, maximum=15, value=5, step=0.1, label="Video Guidance Scale")
-                t2v_generate_btn = gr.Button("Generate Video")
-            with gr.Column():
-                t2v_output = gr.Video(label="Generated Video")
-        
-        t2v_generate_btn.click(
-            generate_video,
-            inputs=[t2v_prompt, t2v_duration, t2v_guidance_scale, t2v_video_guidance_scale],
-            outputs=t2v_output
-        )
+            t2v_generate_btn = gr.Button("Generate Video")
+        with gr.Column():
+            t2v_output = gr.Video(label="Generated Video")
     
-    with gr.Tab("Image-to-Video"):
-        with gr.Row():
-            with gr.Column():
-                i2v_image = gr.Image(type="pil", label="Input Image")
-                i2v_prompt = gr.Textbox(label="Prompt")
-                i2v_duration = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Duration (seconds)")
-                i2v_video_guidance_scale = gr.Slider(minimum=1, maximum=15, value=4, step=0.1, label="Video Guidance Scale")
-                i2v_generate_btn = gr.Button("Generate Video")
-            with gr.Column():
-                i2v_output = gr.Video(label="Generated Video")
+    t2v_generate_btn.click(
+        generate_video,
+        inputs=[i2v_image, t2v_prompt, t2v_duration, t2v_guidance_scale, t2v_video_guidance_scale],
+        outputs=t2v_output
+    )
+    
+    #with gr.Tab("Image-to-Video"):
+    #    with gr.Row():
+    #        with gr.Column():
+                
+    #            i2v_prompt = gr.Textbox(label="Prompt")
+    #            i2v_duration = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Duration (seconds)")
+    #            i2v_video_guidance_scale = gr.Slider(minimum=1, maximum=15, value=4, step=0.1, label="Video Guidance Scale")
+    #            i2v_generate_btn = gr.Button("Generate Video")
+    #        with gr.Column():
+    #            i2v_output = gr.Video(label="Generated Video")
         
-        i2v_generate_btn.click(
-            generate_video_from_image,
-            inputs=[i2v_image, i2v_prompt, i2v_duration, i2v_video_guidance_scale],
-            outputs=i2v_output
-        )
+        #i2v_generate_btn.click(
+        #    generate_video_from_image,
+        #    inputs=[i2v_image, i2v_prompt, i2v_duration, i2v_video_guidance_scale],
+        #    outputs=i2v_output
+        #)
 
 demo.launch()