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()