Create app.py

app.py

@@ -0,0 +1,98 @@
+import gradio as gr
+import torch
+from PIL import Image
+from pyramid_dit import PyramidDiTForVideoGeneration
+from diffusers.utils import load_image, export_to_video
+from huggingface_hub import snapshot_download
+import os
+
+# Download and load the model
+model_path = 'pyramid_flow_model'
+if not os.path.exists(model_path):
+    snapshot_download("rain1011/pyramid-flow-sd3", local_dir=model_path, local_dir_use_symlinks=False, repo_type='model')
+
+torch.cuda.set_device(0)
+model_dtype, torch_dtype = 'bf16', torch.bfloat16
+
+model = PyramidDiTForVideoGeneration(
+    model_path,
+    model_dtype,
+    model_variant='diffusion_transformer_768p',
+)
+
+model.vae.to("cuda")
+model.dit.to("cuda")
+model.text_encoder.to("cuda")
+model.vae.enable_tiling()
+
+def generate_video(prompt, height, width, duration, guidance_scale, video_guidance_scale):
+    temp = 16 if duration == "5s" else 31
+    
+    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=height,
+            width=width,
+            temp=temp,
+            guidance_scale=guidance_scale,
+            video_guidance_scale=video_guidance_scale,
+            output_type="pil",
+        )
+    
+    output_path = "generated_video.mp4"
+    export_to_video(frames, output_path, fps=24)
+    return output_path
+
+def generate_video_from_image(image, prompt, video_guidance_scale):
+    with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
+        frames = model.generate_i2v(
+            prompt=prompt,
+            input_image=image,
+            num_inference_steps=[10, 10, 10],
+            temp=16,
+            video_guidance_scale=video_guidance_scale,
+            output_type="pil",
+        )
+    
+    output_path = "generated_video_from_image.mp4"
+    export_to_video(frames, output_path, fps=24)
+    return output_path
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Pyramid Flow Video Generation Demo")
+    
+    with gr.Tab("Text-to-Video"):
+        with gr.Row():
+            with gr.Column():
+                txt_prompt = gr.Textbox(label="Prompt")
+                txt_height = gr.Slider(384, 768, value=768, step=384, label="Height")
+                txt_width = gr.Slider(640, 1280, value=1280, step=640, label="Width")
+                txt_duration = gr.Radio(["5s", "10s"], value="5s", label="Duration")
+                txt_guidance_scale = gr.Slider(1, 15, value=9, step=0.1, label="Guidance Scale")
+                txt_video_guidance_scale = gr.Slider(1, 15, value=5, step=0.1, label="Video Guidance Scale")
+                txt_generate = gr.Button("Generate Video")
+            with gr.Column():
+                txt_output = gr.Video(label="Generated Video")
+    
+    with gr.Tab("Image-to-Video"):
+        with gr.Row():
+            with gr.Column():
+                img_input = gr.Image(type="pil", label="Input Image")
+                img_prompt = gr.Textbox(label="Prompt (optional)")
+                img_video_guidance_scale = gr.Slider(1, 15, value=4, step=0.1, label="Video Guidance Scale")
+                img_generate = gr.Button("Generate Video")
+            with gr.Column():
+                img_output = gr.Video(label="Generated Video")
+    
+    txt_generate.click(generate_video, 
+                       inputs=[txt_prompt, txt_height, txt_width, txt_duration, txt_guidance_scale, txt_video_guidance_scale], 
+                       outputs=txt_output)
+    
+    img_generate.click(generate_video_from_image,
+                       inputs=[img_input, img_prompt, img_video_guidance_scale],
+                       outputs=img_output)
+
+demo.launch()