Merge pull request #77 from tpc2233/main

Gradio Multi-GPU moving engine scripts to script folder

app_multigpu.py

@@ -14,7 +14,7 @@ def run_inference_multigpu(gpus, variant, model_path, temp, guidance_scale, vide
         output_video = os.path.join(tmpdir, f"{uuid.uuid4()}_output.mp4")
 
         # Path to the external shell script
-        script_path = "./app_multigpu_engine.sh"  # Updated script name
+        script_path = "./scripts/app_multigpu_engine.sh"  # Updated script path
 
         # Prepare the command
         cmd = [
@@ -141,4 +141,3 @@ def generate_text_to_video(prompt, temp, guidance_scale, video_guidance_scale, r
 
 # Launch Gradio app
 demo.launch(share=True)
-

app_multigpu_engine.py → scripts/app_multigpu_engine.py

@@ -1,17 +1,24 @@
 import os
-import torch
 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
-from PIL import Image
 
 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('--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")
@@ -27,7 +34,7 @@ def get_args():
 def main():
     args = get_args()
 
-    # setup DDP
+    # Setup DDP
     init_distributed_mode(args)
 
     assert args.world_size == args.sp_group_size, "The sequence parallel size should match DDP world size"
@@ -68,7 +75,7 @@ def main():
     try:
         if args.task == 't2v':
             prompt = args.prompt
-            with torch.no_grad(), torch.cuda.amp.autocast(enabled=True if model_dtype != 'fp32' else False, dtype=torch_dtype):
+            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],
@@ -94,7 +101,7 @@ def main():
 
             prompt = args.prompt
 
-            with torch.no_grad(), torch.cuda.amp.autocast(enabled=True if model_dtype != 'fp32' else False, dtype=torch_dtype):
+            with torch.no_grad(), torch.cuda.amp.autocast(enabled=(model_dtype != 'fp32'), dtype=torch_dtype):
                 frames = model.generate_i2v(
                     prompt=prompt,
                     input_image=image,

app_multigpu_engine.sh → scripts/app_multigpu_engine.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 
 # Usage:
-# ./app_multigpu_engine.sh GPUS VARIANT MODEL_PATH TASK TEMP GUIDANCE_SCALE VIDEO_GUIDANCE_SCALE RESOLUTION OUTPUT_PATH [IMAGE_PATH] PROMPT
+# ./scripts/app_multigpu_engine.sh GPUS VARIANT MODEL_PATH TASK TEMP GUIDANCE_SCALE VIDEO_GUIDANCE_SCALE RESOLUTION OUTPUT_PATH [IMAGE_PATH] PROMPT
 
 GPUS=$1
 VARIANT=$2
@@ -27,8 +27,19 @@ else
     exit 1
 fi
 
+# Get the directory where the script is located
+SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+# Get the project root directory (parent directory of scripts)
+PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
+
+# Set PYTHONPATH to include the project root directory
+export PYTHONPATH="$PROJECT_ROOT:$PYTHONPATH"
+
+# Adjust the path to app_multigpu_engine.py
+PYTHON_SCRIPT="$SCRIPT_DIR/app_multigpu_engine.py"
+
 torchrun --nproc_per_node="$GPUS" \
-    app_multigpu_engine.py \
+    "$PYTHON_SCRIPT" \
     --model_path "$MODEL_PATH" \
     --variant "$VARIANT" \
     --task "$TASK" \