flash-attn

app.py

@@ -1,7 +1,8 @@
+import os
+import shlex
 import subprocess
 
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
-               shell=True)
+subprocess.run(shlex.split('pip install flash-attn --no-build-isolation'), env=os.environ | {'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"})
 
 import gradio as gr
 import torch
@@ -42,7 +43,7 @@ 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,
@@ -55,7 +56,7 @@ def generate_video(prompt, height, width, duration, guidance_scale, video_guidan
             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
@@ -70,7 +71,7 @@ def generate_video_from_image(image, prompt, video_guidance_scale):
             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
@@ -78,7 +79,7 @@ def generate_video_from_image(image, prompt, video_guidance_scale):
 # 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():
@@ -91,7 +92,7 @@ with gr.Blocks() as demo:
                 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():
@@ -101,13 +102,13 @@ with gr.Blocks() as demo:
                 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], 
+
+    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()
+demo.launch()