Spaces:
Running
Running
File size: 2,570 Bytes
897f552 0f8f3f1 897f552 4e485ec 0a89295 897f552 4e485ec 77adc33 4e485ec 897f552 4e485ec 77adc33 4e485ec 897f552 4e485ec 897f552 3196e8f 77adc33 e20e208 897f552 3196e8f 897f552 4e485ec |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
import argparse
import os
import subprocess # Ensure subprocess is imported
import torch
from huggingface_hub import snapshot_download
# Arguments
parser = argparse.ArgumentParser()
parser.add_argument("--task", type=str, default="t2v-1.3B")
parser.add_argument("--size", type=str, default="200*200")
parser.add_argument("--frame_num", type=int, default=60)
parser.add_argument("--sample_steps", type=int, default=20)
parser.add_argument("--ckpt_dir", type=str, default="./Wan2.1-T2V-1.3B")
parser.add_argument("--offload_model", type=bool, default=True, help="Whether to offload the model (True/False)")
parser.add_argument("--t5_cpu", action="store_true", help="Use CPU for T5 model (optional)")
parser.add_argument("--sample_shift", type=int, default=8, help="Sampling shift for generation")
parser.add_argument("--sample_guide_scale", type=int, default=6, help="Sampling guide scale for generation")
parser.add_argument("--prompt", type=str, required=True)
args = parser.parse_args()
# Log input parameters
print(f"Generating video with the following settings:\n"
f"Task: {args.task}\n"
f"Resolution: {args.size}\n"
f"Frames: {args.frame_num}\n"
f"Sample Steps: {args.sample_steps}\n"
f"Prompt: {args.prompt}\n"
f"Sample Shift: {args.sample_shift}\n"
f"Sample Guide Scale: {args.sample_guide_scale}\n"
f"Using T5 on CPU: {args.t5_cpu}\n"
f"Offload Model: {args.offload_model}")
# Ensure the model is downloaded
if not os.path.exists(args.ckpt_dir):
print("🔄 Downloading WAN 2.1 - 1.3B model from Hugging Face...")
snapshot_download(repo_id="Wan-AI/Wan2.1-T2V-1.3B", local_dir=args.ckpt_dir)
# Free up GPU memory
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
# Command to run the model generation
offload_model_value = "True" if args.offload_model else "False"
command = f"python generate.py --task t2v-1.3B --size 480832 --frame_num 1 --ckpt_dir ./Wan2.1-T2V-1.3B --offload_model True --t5_cpu --sample_shift 8 --sample_guide_scale 6 --prompt \"{args.prompt}\""
# Run the model
process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stdout, stderr = process.communicate()
# Print logs for debugging
print("🔹 Output:", stdout.decode())
print("🔺 Error:", stderr.decode())
# Verify if video was created
if os.path.exists("output.mp4"):
print("✅ Video generated successfully: output.mp4")
else:
print("❌ Error: Video file not found!")
|