Update app.py

app.py

@@ -1,127 +1,143 @@
 import os
-import sys
-import torch
-import argparse
-from PIL import Image
-from diffusers.utils import export_to_video
-
-# Add the project root directory to sys.path
-SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
-PROJECT_ROOT = os.path.dirname(SCRIPT_DIR)
-if PROJECT_ROOT not in sys.path:
-    sys.path.insert(0, PROJECT_ROOT)
-
-from pyramid_dit import PyramidDiTForVideoGeneration
-from trainer_misc import init_distributed_mode, init_sequence_parallel_group
-
-def get_args():
-    parser = argparse.ArgumentParser('Pytorch Multi-process Script', add_help=False)
-    parser.add_argument('--model_dtype', default='bf16', type=str, help="The Model Dtype: bf16")
-    parser.add_argument('--model_path', required=True, type=str, help='Path to the downloaded checkpoint directory')
-    parser.add_argument('--variant', default='diffusion_transformer_768p', type=str)
-    parser.add_argument('--task', default='t2v', type=str, choices=['i2v', 't2v'])
-    parser.add_argument('--temp', default=16, type=int, help='The generated latent num, num_frames = temp * 8 + 1')
-    parser.add_argument('--sp_group_size', default=2, type=int, help="The number of GPUs used for inference, should be 2 or 4")
-    parser.add_argument('--sp_proc_num', default=-1, type=int, help="The number of processes used for video training, default=-1 means using all processes.")
-    parser.add_argument('--prompt', type=str, required=True, help="Text prompt for video generation")
-    parser.add_argument('--image_path', type=str, help="Path to the input image for image-to-video")
-    parser.add_argument('--video_guidance_scale', type=float, default=5.0, help="Video guidance scale")
-    parser.add_argument('--guidance_scale', type=float, default=9.0, help="Guidance scale for text-to-video")
-    parser.add_argument('--resolution', type=str, default='768p', choices=['768p', '384p'], help="Model resolution")
-    parser.add_argument('--output_path', type=str, required=True, help="Path to save the generated video")
-    return parser.parse_args()
-
-def main():
-    args = get_args()
-
-    # Setup DDP
-    init_distributed_mode(args)
-
-    assert args.world_size == args.sp_group_size, "The sequence parallel size should match DDP world size"
-
-    # Enable sequence parallel
-    init_sequence_parallel_group(args)
-
-    device = torch.device('cuda')
-    rank = args.rank
-    model_dtype = args.model_dtype
-
-    model = PyramidDiTForVideoGeneration(
-        args.model_path,
-        model_dtype,
-        model_variant=args.variant,
-    )
-
-    model.vae.to(device)
-    model.dit.to(device)
-    model.text_encoder.to(device)
-    model.vae.enable_tiling()
-
-    if model_dtype == "bf16":
-        torch_dtype = torch.bfloat16 
-    elif model_dtype == "fp16":
-        torch_dtype = torch.float16
+import uuid
+import gradio as gr
+import subprocess
+import tempfile
+import shutil
+
+def run_inference_multigpu(gpus, variant, model_path, temp, guidance_scale, video_guidance_scale, resolution, prompt):
+    """
+    Runs the external multi-GPU inference script and returns the path to the generated video.
+    """
+    # Create a temporary directory to store inputs and outputs
+    with tempfile.TemporaryDirectory() as tmpdir:
+        output_video = os.path.join(tmpdir, f"{uuid.uuid4()}_output.mp4")
+        
+        # Path to the external shell script
+        script_path = "./scripts/app_multigpu_engine.sh"  # Updated script path
+
+        # Prepare the command
+        cmd = [
+            script_path,
+            str(gpus),
+            variant,
+            model_path,
+            't2v',  # Task is always 't2v' since 'i2v' is removed
+            str(temp),
+            str(guidance_scale),
+            str(video_guidance_scale),
+            resolution,
+            output_video,
+            prompt  # Pass the prompt directly as an argument
+        ]
+
+        try:
+            # Run the external script
+            subprocess.run(cmd, check=True)
+        except subprocess.CalledProcessError as e:
+            raise RuntimeError(f"Error during video generation: {e}")
+
+        # After generation, move the video to a permanent location
+        final_output = os.path.join("generated_videos", f"{uuid.uuid4()}_output.mp4")
+        os.makedirs("generated_videos", exist_ok=True)
+        shutil.move(output_video, final_output)
+
+        return final_output
+
+def generate_text_to_video(prompt, temp, guidance_scale, video_guidance_scale, resolution, gpus):
+    model_path = "./pyramid_flow_model"  # Use the model path as specified
+    # Determine variant based on resolution
+    if resolution == "768p":
+        variant = "diffusion_transformer_768p"
     else:
-        torch_dtype = torch.float32
+        variant = "diffusion_transformer_384p"
+    return run_inference_multigpu(gpus, variant, model_path, temp, guidance_scale, video_guidance_scale, resolution, prompt)
 
-    # The video generation config
-    if args.resolution == '768p':
-        width = 1280
-        height = 768
-    else:
-        width = 640
-        height = 384
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+# Pyramid Flow Video Generation Demo
 
-    try:
-        if args.task == 't2v':
-            prompt = args.prompt
-            with torch.no_grad(), torch.cuda.amp.autocast(enabled=(model_dtype != 'fp32'), 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=args.temp,
-                    guidance_scale=args.guidance_scale,
-                    video_guidance_scale=args.video_guidance_scale,
-                    output_type="pil",
-                    save_memory=True,
-                    cpu_offloading=True,
-                    inference_multigpu=True,
-                )
-            if rank == 0:
-                export_to_video(frames, args.output_path, fps=24)
+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.
 
-        elif args.task == 'i2v':
-            if not args.image_path:
-                raise ValueError("Image path is required for image-to-video task")
-            image = Image.open(args.image_path).convert("RGB")
-            image = image.resize((width, height))
-
-            prompt = args.prompt
+[[Paper]](https://arxiv.org/abs/2410.05954) [[Project Page]](https://pyramid-flow.github.io) [[Code]](https://github.com/jy0205/Pyramid-Flow) [[Model]](https://huggingface.co/rain1011/pyramid-flow-sd3)
+"""
+    )
 
-            with torch.no_grad(), torch.cuda.amp.autocast(enabled=(model_dtype != 'fp32'), dtype=torch_dtype):
-                frames = model.generate_i2v(
-                    prompt=prompt,
-                    input_image=image,
-                    num_inference_steps=[10, 10, 10],
-                    temp=args.temp,
-                    video_guidance_scale=args.video_guidance_scale,
-                    output_type="pil",
-                    save_memory=True,
-                    cpu_offloading=True,
-                    inference_multigpu=True,
+    # Shared settings
+    with gr.Row():
+        gpus_dropdown = gr.Dropdown(
+            choices=[2, 4],
+            value=4,
+            label="Number of GPUs"
+        )
+        resolution_dropdown = gr.Dropdown(
+            choices=["768p", "384p"],
+            value="768p",
+            label="Model Resolution"
+        )
+
+    with gr.Tab("Text-to-Video"):
+        with gr.Row():
+            with gr.Column():
+                text_prompt = gr.Textbox(
+                    label="Prompt (Less than 128 words)",
+                    placeholder="Enter a text prompt for the video",
+                    lines=2
                 )
-            if rank == 0:
-                export_to_video(frames, args.output_path, fps=24)
-
-    except Exception as e:
-        if rank == 0:
-            print(f"[ERROR] Error during video generation: {e}")
-        raise
-    finally:
-        torch.distributed.barrier()
+                temp_slider = gr.Slider(1, 31, value=16, step=1, label="Duration")
+                guidance_scale_slider = gr.Slider(1.0, 15.0, value=9.0, step=0.1, label="Guidance Scale")
+                video_guidance_scale_slider = gr.Slider(1.0, 10.0, value=5.0, step=0.1, label="Video Guidance Scale")
+                txt_generate = gr.Button("Generate Video")
+            with gr.Column():
+                txt_output = gr.Video(label="Generated Video")
+        gr.Examples(
+            examples=[
+                [
+                    "A movie trailer featuring the adventures of the 30 year old space man wearing a red wool knitted motorcycle helmet, blue sky, salt desert, cinematic style, shot on 35mm film, vivid colors",
+                    16,
+                    9.0,
+                    5.0,
+                    "768p",
+                    4
+                ],
+                [
+                    "Beautiful, snowy Tokyo city is bustling. The camera moves through the bustling city street, following several people enjoying the beautiful snowy weather and shopping at nearby stalls. Gorgeous sakura petals are flying through the wind along with snowflakes",
+                    16,
+                    9.0,
+                    5.0,
+                    "768p",
+                    4
+                ],
+                [
+                    "Extreme close-up of chicken and green pepper kebabs grilling on a barbeque with flames. Shallow focus and light smoke. vivid colours",
+                    31,
+                    9.0,
+                    5.0,
+                    "768p",
+                    4
+                ],
+            ],
+            inputs=[text_prompt, temp_slider, guidance_scale_slider, video_guidance_scale_slider, resolution_dropdown, gpus_dropdown],
+            outputs=[txt_output],
+            fn=generate_text_to_video,
+            cache_examples='lazy',
+        )
+
+    # Update generate function for Text-to-Video
+    txt_generate.click(
+        generate_text_to_video,
+        inputs=[
+            text_prompt,
+            temp_slider,
+            guidance_scale_slider,
+            video_guidance_scale_slider,
+            resolution_dropdown,
+            gpus_dropdown
+        ],
+        outputs=txt_output
+    )
 
-if __name__ == "__main__":
-    main()
+# Launch Gradio app
+demo.launch(share=False)