mirror of
https://github.com/Wan-Video/Wan2.1.git
synced 2025-11-04 14:16:57 +00:00
Merge pull request #58 from AmericanPresidentJimmyCarter/add-i2v-script
Add a script to generate a single i2v video
This commit is contained in:
deepbeepmeep
GitHub
commit
48b777ae6c
679
i2v_inference.py
Normal file
679
i2v_inference.py
Normal file
@ -0,0 +1,679 @@
|
||||
import os
|
||||
import time
|
||||
import argparse
|
||||
import json
|
||||
import torch
|
||||
import traceback
|
||||
import gc
|
||||
import random
|
||||
|
||||
# These imports rely on your existing code structure
|
||||
# They must match the location of your WAN code, etc.
|
||||
import wan
|
||||
from wan.configs import MAX_AREA_CONFIGS, WAN_CONFIGS
|
||||
from wan.modules.attention import get_attention_modes
|
||||
from wan.utils.utils import cache_video
|
||||
from mmgp import offload, safetensors2, profile_type
|
||||
|
||||
try:
|
||||
import triton
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
DATA_DIR = "ckpts"
|
||||
|
||||
# --------------------------------------------------
|
||||
# HELPER FUNCTIONS
|
||||
# --------------------------------------------------
|
||||
|
||||
def sanitize_file_name(file_name):
|
||||
"""Clean up file name from special chars."""
|
||||
return (
|
||||
file_name.replace("/", "")
|
||||
.replace("\\", "")
|
||||
.replace(":", "")
|
||||
.replace("|", "")
|
||||
.replace("?", "")
|
||||
.replace("<", "")
|
||||
.replace(">", "")
|
||||
.replace('"', "")
|
||||
)
|
||||
|
||||
def extract_preset(lset_name, lora_dir, loras):
|
||||
"""
|
||||
Load a .lset JSON that lists the LoRA files to apply, plus multipliers
|
||||
and possibly a suggested prompt prefix.
|
||||
"""
|
||||
lset_name = sanitize_file_name(lset_name)
|
||||
if not lset_name.endswith(".lset"):
|
||||
lset_name_filename = os.path.join(lora_dir, lset_name + ".lset")
|
||||
else:
|
||||
lset_name_filename = os.path.join(lora_dir, lset_name)
|
||||
|
||||
if not os.path.isfile(lset_name_filename):
|
||||
raise ValueError(f"Preset '{lset_name}' not found in {lora_dir}")
|
||||
|
||||
with open(lset_name_filename, "r", encoding="utf-8") as reader:
|
||||
text = reader.read()
|
||||
lset = json.loads(text)
|
||||
|
||||
loras_choices_files = lset["loras"]
|
||||
loras_choices = []
|
||||
missing_loras = []
|
||||
for lora_file in loras_choices_files:
|
||||
# Build absolute path and see if it is in loras
|
||||
full_lora_path = os.path.join(lora_dir, lora_file)
|
||||
if full_lora_path in loras:
|
||||
idx = loras.index(full_lora_path)
|
||||
loras_choices.append(str(idx))
|
||||
else:
|
||||
missing_loras.append(lora_file)
|
||||
|
||||
if len(missing_loras) > 0:
|
||||
missing_list = ", ".join(missing_loras)
|
||||
raise ValueError(f"Missing LoRA files for preset: {missing_list}")
|
||||
|
||||
loras_mult_choices = lset["loras_mult"]
|
||||
prompt_prefix = lset.get("prompt", "")
|
||||
full_prompt = lset.get("full_prompt", False)
|
||||
return loras_choices, loras_mult_choices, prompt_prefix, full_prompt
|
||||
|
||||
def get_attention_mode(args_attention, installed_modes):
|
||||
"""
|
||||
Decide which attention mode to use: either the user choice or auto fallback.
|
||||
"""
|
||||
if args_attention == "auto":
|
||||
for candidate in ["sage2", "sage", "sdpa"]:
|
||||
if candidate in installed_modes:
|
||||
return candidate
|
||||
return "sdpa" # last fallback
|
||||
elif args_attention in installed_modes:
|
||||
return args_attention
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Requested attention mode '{args_attention}' not installed. "
|
||||
f"Installed modes: {installed_modes}"
|
||||
)
|
||||
|
||||
def load_i2v_model(model_filename, text_encoder_filename, is_720p):
|
||||
"""
|
||||
Load the i2v model with a specific size config and text encoder.
|
||||
"""
|
||||
if is_720p:
|
||||
print("Loading 14B-720p i2v model ...")
|
||||
cfg = WAN_CONFIGS['i2v-14B']
|
||||
wan_model = wan.WanI2V(
|
||||
config=cfg,
|
||||
checkpoint_dir=DATA_DIR,
|
||||
device_id=0,
|
||||
rank=0,
|
||||
t5_fsdp=False,
|
||||
dit_fsdp=False,
|
||||
use_usp=False,
|
||||
i2v720p=True,
|
||||
model_filename=model_filename,
|
||||
text_encoder_filename=text_encoder_filename
|
||||
)
|
||||
else:
|
||||
print("Loading 14B-480p i2v model ...")
|
||||
cfg = WAN_CONFIGS['i2v-14B']
|
||||
wan_model = wan.WanI2V(
|
||||
config=cfg,
|
||||
checkpoint_dir=DATA_DIR,
|
||||
device_id=0,
|
||||
rank=0,
|
||||
t5_fsdp=False,
|
||||
dit_fsdp=False,
|
||||
use_usp=False,
|
||||
i2v720p=False,
|
||||
model_filename=model_filename,
|
||||
text_encoder_filename=text_encoder_filename
|
||||
)
|
||||
# Pipe structure
|
||||
pipe = {
|
||||
"transformer": wan_model.model,
|
||||
"text_encoder": wan_model.text_encoder.model,
|
||||
"text_encoder_2": wan_model.clip.model,
|
||||
"vae": wan_model.vae.model
|
||||
}
|
||||
return wan_model, pipe
|
||||
|
||||
def setup_loras(pipe, lora_dir, lora_preset, num_inference_steps):
|
||||
"""
|
||||
Load loras from a directory, optionally apply a preset.
|
||||
"""
|
||||
from pathlib import Path
|
||||
import glob
|
||||
|
||||
if not lora_dir or not Path(lora_dir).is_dir():
|
||||
print("No valid --lora-dir provided or directory doesn't exist, skipping LoRA setup.")
|
||||
return [], [], [], "", "", False
|
||||
|
||||
# Gather LoRA files
|
||||
loras = sorted(
|
||||
glob.glob(os.path.join(lora_dir, "*.sft"))
|
||||
+ glob.glob(os.path.join(lora_dir, "*.safetensors"))
|
||||
)
|
||||
loras_names = [Path(x).stem for x in loras]
|
||||
|
||||
# Offload them with no activation
|
||||
offload.load_loras_into_model(pipe["transformer"], loras, activate_all_loras=False)
|
||||
|
||||
# If user gave a preset, apply it
|
||||
default_loras_choices = []
|
||||
default_loras_multis_str = ""
|
||||
default_prompt_prefix = ""
|
||||
preset_applied_full_prompt = False
|
||||
if lora_preset:
|
||||
loras_choices, loras_mult, prefix, full_prompt = extract_preset(lora_preset, lora_dir, loras)
|
||||
default_loras_choices = loras_choices
|
||||
# If user stored loras_mult as a list or string in JSON, unify that to str
|
||||
if isinstance(loras_mult, list):
|
||||
# Just store them in a single line
|
||||
default_loras_multis_str = " ".join([str(x) for x in loras_mult])
|
||||
else:
|
||||
default_loras_multis_str = str(loras_mult)
|
||||
default_prompt_prefix = prefix
|
||||
preset_applied_full_prompt = full_prompt
|
||||
|
||||
return (
|
||||
loras,
|
||||
loras_names,
|
||||
default_loras_choices,
|
||||
default_loras_multis_str,
|
||||
default_prompt_prefix,
|
||||
preset_applied_full_prompt
|
||||
)
|
||||
|
||||
def parse_loras_and_activate(
|
||||
transformer,
|
||||
loras,
|
||||
loras_choices,
|
||||
loras_mult_str,
|
||||
num_inference_steps
|
||||
):
|
||||
"""
|
||||
Activate the chosen LoRAs with multipliers over the pipeline's transformer.
|
||||
Supports stepwise expansions (like "0.5,0.8" for partial steps).
|
||||
"""
|
||||
if not loras or not loras_choices:
|
||||
# no LoRAs selected
|
||||
return
|
||||
|
||||
# Handle multipliers
|
||||
def is_float_or_comma_list(x):
|
||||
"""
|
||||
Example: "0.5", or "0.8,1.0", etc. is valid.
|
||||
"""
|
||||
if not x:
|
||||
return False
|
||||
for chunk in x.split(","):
|
||||
try:
|
||||
float(chunk.strip())
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
# Convert multiline or spaced lines to a single list
|
||||
lines = [
|
||||
line.strip()
|
||||
for line in loras_mult_str.replace("\r", "\n").split("\n")
|
||||
if line.strip() and not line.strip().startswith("#")
|
||||
]
|
||||
# Now combine them by space
|
||||
joined_line = " ".join(lines) # "1.0 2.0,3.0"
|
||||
if not joined_line.strip():
|
||||
multipliers = []
|
||||
else:
|
||||
multipliers = joined_line.split(" ")
|
||||
|
||||
# Expand each item
|
||||
final_multipliers = []
|
||||
for mult in multipliers:
|
||||
mult = mult.strip()
|
||||
if not mult:
|
||||
continue
|
||||
if is_float_or_comma_list(mult):
|
||||
# Could be "0.7" or "0.5,0.6"
|
||||
if "," in mult:
|
||||
# expand over steps
|
||||
chunk_vals = [float(x.strip()) for x in mult.split(",")]
|
||||
expanded = expand_list_over_steps(chunk_vals, num_inference_steps)
|
||||
final_multipliers.append(expanded)
|
||||
else:
|
||||
final_multipliers.append(float(mult))
|
||||
else:
|
||||
raise ValueError(f"Invalid LoRA multiplier: '{mult}'")
|
||||
|
||||
# If fewer multipliers than chosen LoRAs => pad with 1.0
|
||||
needed = len(loras_choices) - len(final_multipliers)
|
||||
if needed > 0:
|
||||
final_multipliers += [1.0]*needed
|
||||
|
||||
# Actually activate them
|
||||
offload.activate_loras(transformer, loras_choices, final_multipliers)
|
||||
|
||||
def expand_list_over_steps(short_list, num_steps):
|
||||
"""
|
||||
If user gave (0.5, 0.8) for example, expand them over `num_steps`.
|
||||
The expansion is simply linear slice across steps.
|
||||
"""
|
||||
result = []
|
||||
inc = len(short_list) / float(num_steps)
|
||||
idxf = 0.0
|
||||
for _ in range(num_steps):
|
||||
value = short_list[int(idxf)]
|
||||
result.append(value)
|
||||
idxf += inc
|
||||
return result
|
||||
|
||||
def download_models_if_needed(transformer_filename_i2v, text_encoder_filename, local_folder=DATA_DIR):
|
||||
"""
|
||||
Checks if all required WAN 2.1 i2v files exist locally under 'ckpts/'.
|
||||
If not, downloads them from a Hugging Face Hub repo.
|
||||
Adjust the 'repo_id' and needed files as appropriate.
|
||||
"""
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download, snapshot_download
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"huggingface_hub is required for automatic model download. "
|
||||
"Please install it via `pip install huggingface_hub`."
|
||||
) from e
|
||||
|
||||
# Identify just the filename portion for each path
|
||||
def basename(path_str):
|
||||
return os.path.basename(path_str)
|
||||
|
||||
repo_id = "DeepBeepMeep/Wan2.1"
|
||||
target_root = local_folder
|
||||
|
||||
# You can customize this list as needed for i2v usage.
|
||||
# At minimum you need:
|
||||
# 1) The requested i2v transformer file
|
||||
# 2) The requested text encoder file
|
||||
# 3) VAE file
|
||||
# 4) The open-clip xlm-roberta-large weights
|
||||
#
|
||||
# If your i2v config references additional files, add them here.
|
||||
needed_files = [
|
||||
"Wan2.1_VAE.pth",
|
||||
"models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth",
|
||||
basename(text_encoder_filename),
|
||||
basename(transformer_filename_i2v),
|
||||
]
|
||||
|
||||
# The original script also downloads an entire "xlm-roberta-large" folder
|
||||
# via snapshot_download. If you require that for your pipeline,
|
||||
# you can add it here, for example:
|
||||
subfolder_name = "xlm-roberta-large"
|
||||
if not Path(os.path.join(target_root, subfolder_name)).exists():
|
||||
snapshot_download(repo_id=repo_id, allow_patterns=subfolder_name + "/*", local_dir=target_root)
|
||||
|
||||
for filename in needed_files:
|
||||
local_path = os.path.join(target_root, filename)
|
||||
if not os.path.isfile(local_path):
|
||||
print(f"File '{filename}' not found locally. Downloading from {repo_id} ...")
|
||||
hf_hub_download(
|
||||
repo_id=repo_id,
|
||||
filename=filename,
|
||||
local_dir=target_root
|
||||
)
|
||||
else:
|
||||
# Already present
|
||||
pass
|
||||
|
||||
print("All required i2v files are present.")
|
||||
|
||||
|
||||
# --------------------------------------------------
|
||||
# ARGUMENT PARSER
|
||||
# --------------------------------------------------
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Image-to-Video inference using WAN 2.1 i2v"
|
||||
)
|
||||
# Model + Tools
|
||||
parser.add_argument(
|
||||
"--quantize-transformer",
|
||||
action="store_true",
|
||||
help="Use on-the-fly transformer quantization"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--compile",
|
||||
action="store_true",
|
||||
help="Enable PyTorch 2.0 compile for the transformer"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--attention",
|
||||
type=str,
|
||||
default="auto",
|
||||
help="Which attention to use: auto, sdpa, sage, sage2, flash"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile",
|
||||
type=int,
|
||||
default=4,
|
||||
help="Memory usage profile number [1..5]; see original script or use 2 if you have low VRAM"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--preload",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Megabytes of the diffusion model to preload in VRAM (only used in some profiles)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--verbose",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Verbosity level [0..5]"
|
||||
)
|
||||
|
||||
# i2v Model
|
||||
parser.add_argument(
|
||||
"--transformer-file",
|
||||
type=str,
|
||||
default=f"{DATA_DIR}/wan2.1_image2video_480p_14B_quanto_int8.safetensors",
|
||||
help="Which i2v model to load"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text-encoder-file",
|
||||
type=str,
|
||||
default=f"{DATA_DIR}/models_t5_umt5-xxl-enc-quanto_int8.safetensors",
|
||||
help="Which text encoder to use"
|
||||
)
|
||||
|
||||
# LoRA
|
||||
parser.add_argument(
|
||||
"--lora-dir",
|
||||
type=str,
|
||||
default="",
|
||||
help="Path to a directory containing i2v LoRAs"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lora-preset",
|
||||
type=str,
|
||||
default="",
|
||||
help="A .lset preset name in the lora_dir to auto-apply"
|
||||
)
|
||||
|
||||
# Generation Options
|
||||
parser.add_argument("--prompt", type=str, default=None, required=True, help="Prompt for generation")
|
||||
parser.add_argument("--negative-prompt", type=str, default="", help="Negative prompt")
|
||||
parser.add_argument("--resolution", type=str, default="832x480", help="WxH")
|
||||
parser.add_argument("--frames", type=int, default=64, help="Number of frames (16=1s if fps=16). Must be multiple of 4 +/- 1 in WAN.")
|
||||
parser.add_argument("--steps", type=int, default=30, help="Number of denoising steps.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=5.0, help="Classifier-free guidance scale")
|
||||
parser.add_argument("--flow-shift", type=float, default=3.0, help="Flow shift parameter. Generally 3.0 for 480p, 5.0 for 720p.")
|
||||
parser.add_argument("--riflex", action="store_true", help="Enable RIFLEx for longer videos")
|
||||
parser.add_argument("--teacache", type=float, default=0.25, help="TeaCache multiplier, e.g. 0.5, 2.0, etc.")
|
||||
parser.add_argument("--teacache-start", type=float, default=0.1, help="Teacache start step percentage [0..100]")
|
||||
parser.add_argument("--seed", type=int, default=-1, help="Random seed. -1 means random each time.")
|
||||
|
||||
# LoRA usage
|
||||
parser.add_argument("--loras-choices", type=str, default="", help="Comma-separated list of chosen LoRA indices or preset names to load. Usually you only use the preset.")
|
||||
parser.add_argument("--loras-mult", type=str, default="", help="Multipliers for each chosen LoRA. Example: '1.0 1.2,1.3' etc.")
|
||||
|
||||
# Input
|
||||
parser.add_argument(
|
||||
"--input-image",
|
||||
type=str,
|
||||
default=None,
|
||||
required=True,
|
||||
help="Path to an input image (or multiple)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-file",
|
||||
type=str,
|
||||
default="output.mp4",
|
||||
help="Where to save the resulting video."
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
# --------------------------------------------------
|
||||
# MAIN
|
||||
# --------------------------------------------------
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
# Setup environment
|
||||
offload.default_verboseLevel = args.verbose
|
||||
installed_attn_modes = get_attention_modes()
|
||||
|
||||
# Decide attention
|
||||
chosen_attention = get_attention_mode(args.attention, installed_attn_modes)
|
||||
offload.shared_state["_attention"] = chosen_attention
|
||||
|
||||
# Determine i2v resolution format
|
||||
if "720" in args.transformer_file:
|
||||
is_720p = True
|
||||
else:
|
||||
is_720p = False
|
||||
|
||||
# Make sure we have the needed models locally
|
||||
download_models_if_needed(args.transformer_file, args.text_encoder_file)
|
||||
|
||||
# Load i2v
|
||||
wan_model, pipe = load_i2v_model(
|
||||
model_filename=args.transformer_file,
|
||||
text_encoder_filename=args.text_encoder_file,
|
||||
is_720p=is_720p
|
||||
)
|
||||
wan_model._interrupt = False
|
||||
|
||||
# Offload / profile
|
||||
# e.g. for your script: offload.profile(pipe, profile_no=args.profile, compile=..., quantizeTransformer=...)
|
||||
# pass the budgets if you want, etc.
|
||||
kwargs = {}
|
||||
if args.profile == 2 or args.profile == 4:
|
||||
# preload is in MB
|
||||
if args.preload == 0:
|
||||
budgets = {"transformer": 100, "text_encoder": 100, "*": 1000}
|
||||
else:
|
||||
budgets = {"transformer": args.preload, "text_encoder": 100, "*": 1000}
|
||||
kwargs["budgets"] = budgets
|
||||
elif args.profile == 3:
|
||||
kwargs["budgets"] = {"*": "70%"}
|
||||
|
||||
compile_choice = "transformer" if args.compile else ""
|
||||
# Create the offload object
|
||||
offloadobj = offload.profile(
|
||||
pipe,
|
||||
profile_no=args.profile,
|
||||
compile=compile_choice,
|
||||
quantizeTransformer=args.quantize_transformer,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
# If user wants to use LoRAs
|
||||
(
|
||||
loras,
|
||||
loras_names,
|
||||
default_loras_choices,
|
||||
default_loras_multis_str,
|
||||
preset_prompt_prefix,
|
||||
preset_full_prompt
|
||||
) = setup_loras(pipe, args.lora_dir, args.lora_preset, args.steps)
|
||||
|
||||
# Combine user prompt with preset prompt if the preset indicates so
|
||||
if preset_prompt_prefix:
|
||||
if preset_full_prompt:
|
||||
# Full override
|
||||
user_prompt = preset_prompt_prefix
|
||||
else:
|
||||
# Just prefix
|
||||
user_prompt = preset_prompt_prefix + "\n" + args.prompt
|
||||
else:
|
||||
user_prompt = args.prompt
|
||||
|
||||
# Actually parse user LoRA choices if they did not rely purely on the preset
|
||||
if args.loras_choices:
|
||||
# If user gave e.g. "0,1", we treat that as new additions
|
||||
lora_choice_list = [x.strip() for x in args.loras_choices.split(",")]
|
||||
else:
|
||||
# Use the defaults from the preset
|
||||
lora_choice_list = default_loras_choices
|
||||
|
||||
# Activate them
|
||||
parse_loras_and_activate(
|
||||
pipe["transformer"], loras, lora_choice_list, args.loras_mult or default_loras_multis_str, args.steps
|
||||
)
|
||||
|
||||
# Negative prompt
|
||||
negative_prompt = args.negative_prompt or ""
|
||||
|
||||
# Sanity check resolution
|
||||
if "*" in args.resolution.lower():
|
||||
print("ERROR: resolution must be e.g. 832x480 not '832*480'. Fixing it.")
|
||||
resolution_str = args.resolution.lower().replace("*", "x")
|
||||
else:
|
||||
resolution_str = args.resolution
|
||||
|
||||
try:
|
||||
width, height = [int(x) for x in resolution_str.split("x")]
|
||||
except:
|
||||
raise ValueError(f"Invalid resolution: '{resolution_str}'")
|
||||
|
||||
# Additional checks (from your original code).
|
||||
if "480p" in args.transformer_file:
|
||||
# Then we cannot exceed certain area for 480p model
|
||||
if width * height > 832*480:
|
||||
raise ValueError("You must use the 720p i2v model to generate bigger than 832x480.")
|
||||
# etc.
|
||||
|
||||
# Handle random seed
|
||||
if args.seed < 0:
|
||||
args.seed = random.randint(0, 999999999)
|
||||
print(f"Using seed={args.seed}")
|
||||
|
||||
# Setup tea cache if needed
|
||||
trans = wan_model.model
|
||||
trans.enable_teacache = (args.teacache > 0)
|
||||
if trans.enable_teacache:
|
||||
if "480p" in args.transformer_file:
|
||||
# example from your code
|
||||
trans.coefficients = [-3.02331670e+02, 2.23948934e+02, -5.25463970e+01, 5.87348440e+00, -2.01973289e-01]
|
||||
elif "720p" in args.transformer_file:
|
||||
trans.coefficients = [-114.36346466, 65.26524496, -18.82220707, 4.91518089, -0.23412683]
|
||||
else:
|
||||
raise ValueError("Teacache not supported for this model variant")
|
||||
|
||||
# Attempt generation
|
||||
print("Starting generation ...")
|
||||
start_time = time.time()
|
||||
|
||||
# Read the input image
|
||||
if not os.path.isfile(args.input_image):
|
||||
raise ValueError(f"Input image does not exist: {args.input_image}")
|
||||
|
||||
from PIL import Image
|
||||
input_img = Image.open(args.input_image).convert("RGB")
|
||||
|
||||
# Possibly load more than one image if you want "multiple images" – but here we'll just do single for demonstration
|
||||
|
||||
# Define the generation call
|
||||
# - frames => must be multiple of 4 plus 1 as per original script's note, e.g. 81, 65, ...
|
||||
# You can correct to that if needed:
|
||||
frame_count = (args.frames // 4)*4 + 1 # ensures it's 4*N+1
|
||||
# RIFLEx
|
||||
enable_riflex = args.riflex
|
||||
|
||||
# If teacache => reset counters
|
||||
if trans.enable_teacache:
|
||||
trans.teacache_counter = 0
|
||||
trans.teacache_multiplier = args.teacache
|
||||
trans.teacache_start_step = int(args.teacache_start * args.steps / 100.0)
|
||||
trans.num_steps = args.steps
|
||||
trans.teacache_skipped_steps = 0
|
||||
trans.previous_residual_uncond = None
|
||||
trans.previous_residual_cond = None
|
||||
|
||||
# VAE Tiling
|
||||
device_mem_capacity = torch.cuda.get_device_properties(0).total_memory / 1048576
|
||||
if device_mem_capacity >= 28000: # 81 frames 720p requires about 28 GB VRAM
|
||||
use_vae_config = 1
|
||||
elif device_mem_capacity >= 8000:
|
||||
use_vae_config = 2
|
||||
else:
|
||||
use_vae_config = 3
|
||||
|
||||
if use_vae_config == 1:
|
||||
VAE_tile_size = 0
|
||||
elif use_vae_config == 2:
|
||||
VAE_tile_size = 256
|
||||
else:
|
||||
VAE_tile_size = 128
|
||||
|
||||
print('Using VAE tile size of', VAE_tile_size)
|
||||
|
||||
# Actually run the i2v generation
|
||||
try:
|
||||
sample_frames = wan_model.generate(
|
||||
user_prompt,
|
||||
input_img,
|
||||
frame_num=frame_count,
|
||||
max_area=MAX_AREA_CONFIGS[f"{width}*{height}"], # or you can pass your custom
|
||||
shift=args.flow_shift,
|
||||
sampling_steps=args.steps,
|
||||
guide_scale=args.guidance_scale,
|
||||
n_prompt=negative_prompt,
|
||||
seed=args.seed,
|
||||
offload_model=False,
|
||||
callback=None, # or define your own callback if you want
|
||||
enable_RIFLEx=enable_riflex,
|
||||
VAE_tile_size=VAE_tile_size,
|
||||
)
|
||||
except Exception as e:
|
||||
offloadobj.unload_all()
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
err_str = f"Generation failed with error: {e}"
|
||||
# Attempt to detect OOM errors
|
||||
s = str(e).lower()
|
||||
if any(keyword in s for keyword in ["memory", "cuda", "alloc"]):
|
||||
raise RuntimeError("Likely out-of-VRAM or out-of-RAM error. " + err_str)
|
||||
else:
|
||||
traceback.print_exc()
|
||||
raise RuntimeError(err_str)
|
||||
|
||||
# After generation
|
||||
offloadobj.unload_all()
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if sample_frames is None:
|
||||
raise RuntimeError("No frames were returned (maybe generation was aborted or failed).")
|
||||
|
||||
# If teacache was used, we can see how many steps were skipped
|
||||
if trans.enable_teacache:
|
||||
print(f"TeaCache skipped steps: {trans.teacache_skipped_steps} / {args.steps}")
|
||||
|
||||
# Save result
|
||||
sample_frames = sample_frames.cpu() # shape = c, t, h, w => [3, T, H, W]
|
||||
os.makedirs(os.path.dirname(args.output_file) or ".", exist_ok=True)
|
||||
|
||||
# Use the provided helper from your code to store the MP4
|
||||
# By default, you used cache_video(tensor=..., save_file=..., fps=16, ...)
|
||||
# or you can do your own. We'll do the same for consistency:
|
||||
cache_video(
|
||||
tensor=sample_frames[None], # shape => [1, c, T, H, W]
|
||||
save_file=args.output_file,
|
||||
fps=16,
|
||||
nrow=1,
|
||||
normalize=True,
|
||||
value_range=(-1, 1)
|
||||
)
|
||||
|
||||
end_time = time.time()
|
||||
elapsed_s = end_time - start_time
|
||||
print(f"Done! Output written to {args.output_file}. Generation time: {elapsed_s:.1f} seconds.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Reference in New Issue
Block a user