Update app.py

app.py

@@ -10,12 +10,15 @@ 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)
+
 is_canonical = True if os.environ.get("SPACE_ID") == "Pyramid-Flow/pyramid-flow" else False
 
 # Constants
 MODEL_PATH = "pyramid-flow-model"
 MODEL_REPO = "rain1011/pyramid-flow-sd3"
-MODEL_VARIANT = "diffusion_transformer_384p"
+MODEL_VARIANT = "diffusion_transformer_768p"
 MODEL_DTYPE = "bf16"
 
 def center_crop(image, target_width, target_height):
@@ -63,18 +66,19 @@ model = load_model()
 # Text-to-video generation function
 @spaces.GPU(duration=120)
 def generate_video(prompt, image=None, duration=5, guidance_scale=9, video_guidance_scale=5, progress=gr.Progress(track_tqdm=True)):
-    multiplier = 3
-    temp = int(duration * multiplier) + 1  # Convert seconds to temp value (assuming 24 FPS)
+    multiplier = 0.8 if is_canonical else 2.4
+    temp = int(duration * 0.8)  # Convert seconds to temp value (assuming 24 FPS)
     torch_dtype = torch.bfloat16 if MODEL_DTYPE == "bf16" else torch.float32
     if(image):
-        cropped_image = center_crop(image, 640, 384)
-        resized_image = cropped_image.resize((640, 384))
+        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,
@@ -85,8 +89,8 @@ def generate_video(prompt, image=None, duration=5, guidance_scale=9, video_guida
                 prompt=prompt,
                 num_inference_steps=[20, 20, 20],
                 video_num_inference_steps=[10, 10, 10],
-                height=384,
-                width=640,
+                height=768,
+                width=1280,
                 temp=temp,
                 guidance_scale=guidance_scale,
                 video_guidance_scale=video_guidance_scale,
@@ -94,14 +98,14 @@ def generate_video(prompt, image=None, duration=5, guidance_scale=9, video_guida
                 save_memory=True,
             )
     output_path = f"{str(uuid.uuid4())}_output_video.mp4"
-    export_to_video(frames, output_path, fps=24)
+    export_to_video(frames, output_path, fps=8 if is_canonical else 24)
     return output_path
 
 # Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Pyramid Flow 384p demo")
-    gr.Markdown("Pyramid Flow is a training-efficient **Autoregressive Video Generation** model based on **Flow Matching**. It is trained only on open-source datasets within 20.7k A100 GPU hours")
-    gr.Markdown("[[Paper](https://arxiv.org/pdf/2410.05954)], [[Model](https://huggingface.co/rain1011/pyramid-flow-sd3)], [[Code](https://github.com/jy0205/Pyramid-Flow)] [[Project Page]](https://pyramid-flow.github.io)")
+    gr.Markdown("# Pyramid Flow")
+    gr.Markdown("Pyramid Flow is a training-efficient Autoregressive Video Generation model based on Flow Matching. It is trained only on open-source datasets within 20.7k A100 GPU hours")
+    gr.Markdown("[[Paper](https://arxiv.org/pdf/2410.05954)], [[Model](https://huggingface.co/rain1011/pyramid-flow-sd3)], [[Code](https://github.com/jy0205/Pyramid-Flow)]")
     
     with gr.Row():
         with gr.Column():
@@ -109,8 +113,8 @@ with gr.Blocks() as demo:
                 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=5, value=5, step=1, label="Duration (seconds)", visible=not is_canonical)
-                t2v_guidance_scale = gr.Slider(minimum=1, maximum=15, value=7, step=0.1, label="Guidance Scale")
+                t2v_duration = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Duration (seconds)", visible=not is_canonical)
+                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():
@@ -122,7 +126,7 @@ with gr.Blocks() as demo:
                             <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-lg.svg" alt="Duplicate this Space">
                         </a>
                     </p>
-                    <p>to use privately and generate videos up to 10s</p>
+                    <p>to use privately and generate videos up to 10s at 24fps</p>
                 </div>
                 """)
     gr.Examples(