import os import sys import argparse import json import subprocess import torch import multiprocessing as mp import csv import re from pathlib import Path import logging from datetime import datetime import yaml def load_config(config_path): _, ext = os.path.splitext(config_path) with open(config_path, "r") as f: if ext == ".json": return json.load(f) if ext in [".yaml", ".yml"]: return yaml.safe_load(f) raise ValueError(f"Unsupported config file extension: {ext}") def setup_gpu_logger(gpu_id: int, base_log_dir: str, run_timestamp: str): """Set up independent log file for each GPU""" # If base_log_dir contains /videos/ subpath, extract its parent directory if "videos/" in base_log_dir or "/videos/" in base_log_dir: # e.g.: path/output/videos/overall_consistency -> path/output if "/videos/" in base_log_dir: base_log_dir = base_log_dir.split("/videos/")[0] else: base_log_dir = base_log_dir.split("videos/")[0] # Create subfolder by run time log_dir = os.path.join(base_log_dir, "log", run_timestamp) os.makedirs(log_dir, exist_ok=True) # Simplified log file name log_file = os.path.join(log_dir, f"gpu_{gpu_id}.log") # Create logger logger = logging.getLogger(f"GPU_{gpu_id}") logger.setLevel(logging.INFO) # Clear existing handlers logger.handlers.clear() # File handler file_handler = logging.FileHandler(log_file) file_handler.setLevel(logging.INFO) # Console handler (optional, if you still want to see output in terminal) # console_handler = logging.StreamHandler() # console_handler.setLevel(logging.INFO) # Format formatter = logging.Formatter( f'[GPU {gpu_id}] %(asctime)s - %(levelname)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S' ) file_handler.setFormatter(formatter) # console_handler.setFormatter(formatter) logger.addHandler(file_handler) # logger.addHandler(console_handler) # Also redirect print to log file sys.stdout = LoggerWriter(logger, logging.INFO) sys.stderr = LoggerWriter(logger, logging.ERROR) return logger, log_file class LoggerWriter: """Redirect print output to logger""" def __init__(self, logger, level): self.logger = logger self.level = level def write(self, message): if message.strip(): self.logger.log(self.level, message.rstrip()) def flush(self): pass def worker_process(gpu_id: int, rank: int, args: argparse.Namespace, all_samples: list, run_timestamp: str): """Independent worker running on each GPU""" # ========== For pdb debugging in terminal ========== sys.stdin = open(0) # ================================================== os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) os.environ["WORLD_SIZE"] = "1" os.environ["RANK"] = "0" os.environ["LOCAL_RANK"] = "0" os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = str(29510 + gpu_id) # ========== Set up independent log file for each GPU ========== logger, log_file = setup_gpu_logger(gpu_id, args.save_path, run_timestamp) logger.info(f"GPU {gpu_id} logger initialized. Log file: {log_file}") if getattr(args, "print_peak_memory", False) and torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() device = torch.cuda.current_device() print(f"[GPU {gpu_id}] GPU Memory before pipeline: {torch.cuda.memory_allocated(device) / 1024**3:.2f} GB allocated") from inference.pipeline.video_generate import SampleTransport if args.reuse_strategy == "original": pass elif args.reuse_strategy == "all": from inference.pipeline.teacache_all import teacache_forward_velocity, teacache_integrate_velocity # TeaCache SampleTransport.rel_l1_thresh = args.rel_l1_thresh SampleTransport.accumulated_rel_l1_distance = 0 SampleTransport.previous_modulated_input = None SampleTransport.previous_residual = None SampleTransport.cnt = 0 SampleTransport.forward_velocity = teacache_forward_velocity SampleTransport.integrate_velocity = teacache_integrate_velocity SampleTransport.reuse_times = 0 SampleTransport.no_reuse_first_n_steps = args.no_reuse_first_n_steps SampleTransport.previous_output = None SampleTransport.whole_calc_when_cross = args.whole_calc_when_cross SampleTransport.log = args.log elif args.reuse_strategy == "chunkwise": from inference.pipeline.flowcache import teacache_forward_velocity, teacache_integrate_velocity if args.compress_kv_cache: print("KV cache compression is enabled.") # KV cache compression SampleTransport.compress_kv_cache = True assert args.total_cache_chunk_nums is not None compression_config = { "method_config": { "compress_strategy": args.compress_strategy, "mix_lambda": args.mix_lambda, "query_granularity": args.query_granularity, "score_weighting_method": args.score_weighting_method, "power": args.power, }, } from inference.rkv import replace_magi from inference.rkv.utils import ChunkKVRangeTracker replace_magi(compression_config) else: SampleTransport.compress_kv_cache = False SampleTransport.rel_l1_thresh = args.rel_l1_thresh SampleTransport.chunk_accumulated_rel_l1 = 0 SampleTransport.previous_modulated_input = None SampleTransport.previous_residual = None SampleTransport.cnt = 0 SampleTransport.forward_velocity = teacache_forward_velocity SampleTransport.integrate_velocity = teacache_integrate_velocity SampleTransport.reuse_times = 0 SampleTransport.no_reuse_first_n_steps = args.no_reuse_first_n_steps SampleTransport.previous_output = None SampleTransport.discard_nearly_clean_chunk = args.discard_nearly_clean_chunk SampleTransport.whole_calc_when_cross = args.whole_calc_when_cross SampleTransport.no_reuse_mode = args.no_reuse_mode # Token-wise reuse parameters SampleTransport.token_wise_reuse = args.token_wise_reuse SampleTransport.token_rel_l1_thresh = args.token_rel_l1_thresh if args.token_rel_l1_thresh is not None else args.rel_l1_thresh SampleTransport.tokenwise_l1_mode = args.tokenwise_l1_mode SampleTransport.compute_l1_weights_once = getattr(args, "compute_l1_weights_once", False) # Three-phase control parameters SampleTransport.warmup_steps = args.warmup_steps SampleTransport.chunk_wise_only_steps = args.chunk_wise_only_steps # Token reuse ratio control SampleTransport.max_token_reuse_ratio = args.max_token_reuse_ratio SampleTransport.initial_token_reuse_ratio = args.initial_token_reuse_ratio SampleTransport.final_token_reuse_ratio = args.final_token_reuse_ratio # Temporal weight parameters SampleTransport.temporal_weight_floor = args.temporal_weight_floor SampleTransport.temporal_weight_power = args.temporal_weight_power SampleTransport.enable_temporal_voting = args.enable_temporal_voting # Debug and visualization parameters SampleTransport.print_token_stats = args.print_token_stats SampleTransport.visualize_reuse_mask = args.visualize_reuse_mask SampleTransport.debug = getattr(args, "debug", False) # --- Per-chunk state --- SampleTransport.chunk_accumulated_rel_l1 = None # List[float]: Accumulated rel L1 for each chunk SampleTransport.prev_chunk_features = None # List[Tensor]: Features of each chunk from previous step SampleTransport.chunk_reuse_flags = None # Whether each chunk is reused in current step # KV cache compression SampleTransport.total_cache_chunk_nums = args.total_cache_chunk_nums # Check mutual exclusivity of token_wise_reuse and compress_kv_cache if SampleTransport.token_wise_reuse and SampleTransport.compress_kv_cache: raise ValueError( "token_wise_reuse and compress_kv_cache cannot be enabled simultaneously. " "Token-level reuse requires full chunk query states for KV cache reassembly, " "which is incompatible with KV cache compression. " "Please set only one of these options to True." ) # log SampleTransport.log = args.log else: raise ValueError(f"Unknown reuse strategy: {args.reuse_strategy}") try: magi_root = subprocess.check_output(['git', 'rev-parse', '--show-toplevel']).decode().strip() os.environ["MAGI_ROOT"] = magi_root os.environ["PYTHONPATH"] = f"{magi_root}:{os.environ.get('PYTHONPATH', '')}" except Exception as e: print(f"[GPU {gpu_id}] Failed to set MAGI_ROOT: {e}") return filtered_samples = [] for sample in all_samples: if args.benchmark == 'vbench': prompt = sample safe_prompt = prompt seed = 0 cur_save_path = os.path.join(args.save_path, f"{safe_prompt}-{seed}.mp4") else: # physicsiq cur_save_path = sample['output_path'] cur_save_path = os.path.abspath(cur_save_path) if not os.path.exists(cur_save_path): filtered_samples.append(sample) else: print(f"[✅ SKIP] Already exists: {cur_save_path}") if not filtered_samples: print(f"[GPU {gpu_id}] No samples need to be generated.") return print(f"Totally Processing {len(filtered_samples)} samples.") # === Split samples === world_size = args.num_gpus samples_per_gpu = (len(filtered_samples) + world_size - 1) // world_size start_idx = rank * samples_per_gpu end_idx = min(start_idx + samples_per_gpu, len(filtered_samples)) my_samples = filtered_samples[start_idx:end_idx] if not my_samples: print(f"[GPU {gpu_id}] No samples assigned.") return print(f"[GPU {gpu_id}] Processing {len(my_samples)} samples: {start_idx} ~ {end_idx-1}") from inference.pipeline.entry import MagiPipeline from inference.common import set_random_seed print(f"[GPU {gpu_id}] Loading model...") pipeline = MagiPipeline(args.config_file) print(f"[GPU {gpu_id}] Model loaded successfully.") for i, sample in enumerate(my_samples): if args.benchmark == 'vbench': prompt = sample seed = 0 safe_prompt = prompt output_path = os.path.join(args.save_path, f"{safe_prompt}-{seed}.mp4") output_path = os.path.abspath(output_path) if os.path.exists(output_path): print(f"[✅ GPU {gpu_id}] Already exists: {output_path}") continue # Reset seed before each prompt to ensure reproducibility set_random_seed(pipeline.config.runtime_config.seed) print(f"[🚀 GPU {gpu_id}] Generating T2V: '{prompt}' -> {output_path}") pipeline.run_text_to_video(prompt=prompt, output_path=output_path) print(f"[✅ GPU {gpu_id}] Saved: {output_path}") else: # physicsiq prompt = sample['description'] prefix_video_path = sample['prefix_video_path'] output_path = sample['output_path'] if not os.path.exists(prefix_video_path): print(f"[⚠️ GPU {gpu_id}] Conditioning video not found: {prefix_video_path}") continue if os.path.exists(output_path): print(f"[✅ GPU {gpu_id}] Already exists: {output_path}") continue print(f"[🚀 GPU {gpu_id}] Generating V2V: '{prompt}'") print(f" → Input: {prefix_video_path}") print(f" → Output: {output_path}") # Reset seed before each prompt to ensure reproducibility set_random_seed(pipeline.config.runtime_config.seed) pipeline.run_video_to_video( prompt=prompt, prefix_video_path=prefix_video_path, output_path=output_path ) print(f"[✅ GPU {gpu_id}] Saved: {output_path}") if getattr(args, "print_peak_memory", False) and torch.cuda.is_available(): device = torch.cuda.current_device() print(f"[GPU {gpu_id}] Peak GPU memory: {torch.cuda.max_memory_allocated(device) / 1024**3:.2f} GB") print(f"[GPU {gpu_id}] Done.") def load_physicsiq_samples(args): """Load sample list for PhysicsIQ dataset""" DATA_ROOT = args.physicsiq_data_dir DESCRIPTIONS_CSV = os.path.join(DATA_ROOT, "descriptions", "descriptions.csv") OUTPUT_DIR = args.save_path FPS = 24 # MAGI uses 24 FPS if not os.path.exists(DESCRIPTIONS_CSV): raise FileNotFoundError(f"descriptions.csv not found at {DESCRIPTIONS_CSV}") os.makedirs(OUTPUT_DIR, exist_ok=True) samples = [] with open(DESCRIPTIONS_CSV, mode='r') as f: reader = csv.DictReader(f) for row in reader: scenario = row['scenario'].strip() match = re.match(r'^(\d+)_', scenario) if not match: print(f"Cannot extract ID from scenario: {scenario}") continue vid_id = match.group(1).zfill(4) description = row['description'] generated_video_name = row['generated_video_name'] # Construct conditioning video path conditioning_dir = os.path.join(DATA_ROOT, "physics-IQ-benchmark", "split-videos", "conditioning", f"{FPS}FPS") match_suffix = re.search(r'_(.*)', scenario) prefix_video_name = match_suffix.group(1) if match_suffix else "" filename = f"{vid_id}_conditioning-videos_{FPS}FPS_{prefix_video_name}" prefix_video_path = os.path.join(conditioning_dir, filename) output_path = os.path.join(OUTPUT_DIR, generated_video_name) os.makedirs(os.path.dirname(output_path), exist_ok=True) samples.append({ 'vid_id': vid_id, 'scenario': scenario, 'description': description, 'generated_video_name': generated_video_name, 'prefix_video_path': prefix_video_path, 'output_path': output_path }) total_samples_num = len(samples) unique_samples_num = total_samples_num // 2 print(f"[✅] Loaded {unique_samples_num} samples.") # physicsiq samples twice, we only take the first half filtered_samples = samples[:unique_samples_num] # Apply start/end slice if args.start is not None or args.end is not None: start = args.start if args.start is not None else 0 end = args.end if args.end is not None else len(filtered_samples) # Ensure index is within valid range start = max(0, min(start, len(filtered_samples))) end = max(start, min(end, len(filtered_samples))) filtered_samples = filtered_samples[start:end] print(f"[✂️] Sliced to samples {start} to {end-1}, total {len(filtered_samples)} samples.") return filtered_samples def main(): parser = argparse.ArgumentParser() parser.add_argument("--benchmark", type=str, required=True, choices=['vbench', 'physicsiq'], help="Benchmark name") parser.add_argument("--vbench_prompt_dir", type=str, help="Directory containing prompt files") parser.add_argument("--dimension", type=str, help="Benchmark dimension") parser.add_argument('--physicsiq_data_dir', type=str, default='physics-iq-benchmark', help='Root path of the dataset') parser.add_argument("--save_path", type=str, default="./benchmark_videos", help="Output directory") parser.add_argument("--config_file", type=str, default="example/4.5B/4.5B_distill_config.json", help="Model config file") parser.add_argument("--additional_config", type=str, help="Additional MotionCache config file") parser.add_argument("--gpus", type=str, default="0", help="Comma-separated GPU IDs to use") parser.add_argument("--reuse_strategy", type=str, choices=['all', 'chunkwise', 'original'], required=True, help="Reuse strategy for video generation") parser.add_argument('--rel_l1_thresh', type=float, default=0, help='Relative L1 distance threshold for TeaCache.') parser.add_argument('--no_reuse_first_n_steps', type=int, default=0, help='Number of steps to not reuse output.') # chunkwise reuse parser.add_argument('--discard_nearly_clean_chunk', action='store_true', help='Whether to discard nearly clean chunks.') parser.add_argument('--whole_calc_when_cross', action='store_true', help='Whether to perform whole calculation when crossing stage which number of chunks increases.') parser.add_argument( '--no_reuse_mode', type=str, choices=['first', 'mid', 'none'], help='Detemine which part of steps to not reuse output. first: first n steps; mid: mid n steps; none: not set .' ) # KV cache compression parser.add_argument('--compress_kv_cache', action='store_true', help='Whether to compress kv cache.') parser.add_argument('--total_cache_chunk_nums', type=int, default=None, help='Total length of kv cache.') parser.add_argument('--compress_strategy', type=str, default='token', choices=['token', 'frame', 'chunk'], help='KV cache compression strategy.') parser.add_argument('--mix_lambda', type=float, default=0.07, help='Mix lambda for token compression.') parser.add_argument('--query_granularity', type=str, default='chunk', choices=['chunk', 'frame', 'token'], help='Query granularity for token compression.') parser.add_argument('--score_weighting_method', type=str, default='no_weight', choices=['no_weight','hard_code', 'exponential', 'polynomial', 'gaussian', "upper_convex_polynomial"], help='Score weighting method for KV compression.') parser.add_argument('--power', type=float, default=3, help='Power parameter for polynomial and upper_convex_polynomial weighting methods.') # log parser.add_argument('--log', action='store_true', help='Whether to log the TeaCache information.') parser.add_argument('--print_peak_memory', action='store_true', help='Print peak GPU memory usage.') # Token-wise reuse parameters parser.add_argument('--token_wise_reuse', action='store_true', help='Enable token-wise reuse strategy.') parser.add_argument('--token_rel_l1_thresh', type=float, default=None, help='Token-level relative L1 distance threshold (defaults to rel_l1_thresh if not set).') parser.add_argument('--tokenwise_l1_mode', type=str, default='chunk', choices=['chunk', 'token'], help='Mode for token-level rel_l1 calculation: "chunk" uses chunk-level rel_l1 for all tokens, "token" computes per-token rel_l1 independently.') # Three-phase control parameters parser.add_argument('--warmup_steps', type=int, default=5, help='Number of warmup steps where no caching is used.') parser.add_argument('--chunk_wise_only_steps', type=int, default=0, help='Number of steps for chunk-wise only phase (before token-wise phase).') # Token reuse ratio control parser.add_argument('--max_token_reuse_ratio', type=float, default=1.0, help='Fixed maximum ratio of tokens that can be reused (0~1). Set to 1.0 to allow all tokens to be reused.') parser.add_argument('--initial_token_reuse_ratio', type=float, default=None, help='Initial token reuse ratio at the start of hierarchical phase (for dynamic ratio mode).') parser.add_argument('--final_token_reuse_ratio', type=float, default=None, help='Final token reuse ratio at the end of hierarchical phase (for dynamic ratio mode).') # Temporal weight parameters parser.add_argument('--temporal_weight_floor', type=float, default=0.0, help='Floor for temporal weight normalization, maps weights to [floor, 1] range.') parser.add_argument('--temporal_weight_power', type=float, default=None, help='Power for nonlinear temporal weight normalization (default: None=linear). Values < 1 make more tokens closer to 1 (convex curve), > 1 make more tokens closer to floor.') parser.add_argument('--enable_temporal_voting', action='store_true', help='Enable temporal voting: force tokens at the same spatial position across frames to have the same reuse decision via majority voting.') # Debug and visualization parser.add_argument('--print_token_stats', action='store_true', help='Print token-wise reuse statistics.') parser.add_argument('--visualize_reuse_mask', action='store_true', help='Visualize token reuse mask on the output video.') # sampling range control parser.add_argument('--start', type=int, default=None, help='Start index of samples to process (inclusive)') parser.add_argument('--end', type=int, default=None, help='End index of samples to process (exclusive)') args = parser.parse_args() if args.additional_config: additional_config = load_config(args.additional_config) print(f"Loading additional config: {additional_config}") for key, value in additional_config.items(): setattr(args, key, value) print(f"[✅] Parsed arguments: {args}") if args.benchmark == 'vbench': assert os.path.exists(args.vbench_prompt_dir), f"Prompt directory not found: {args.vbench_prompt_dir}" prompt_file = os.path.join(args.vbench_prompt_dir, f"{args.dimension}.txt") if not os.path.exists(prompt_file): raise FileNotFoundError(f"Prompt file not found: {prompt_file}") with open(prompt_file, 'r') as f: all_samples = [line.strip() for line in f if line.strip()] # Apply start/end slice if args.start is not None or args.end is not None: start = args.start if args.start is not None else 0 end = args.end if args.end is not None else len(all_samples) # Ensure index is within valid range start = max(0, min(start, len(all_samples))) end = max(start, min(end, len(all_samples))) all_samples = all_samples[start:end] print(f"[✂️] Sliced to prompts {start} to {end-1}, total {len(all_samples)} prompts.") elif args.benchmark == 'physicsiq': assert os.path.exists(args.physicsiq_data_dir), f"Data root directory not found: {args.physicsiq_data_dir}" all_samples = load_physicsiq_samples(args) else: raise ValueError(f"Invalid benchmark: {args.benchmark}") gpu_ids = list(map(int, args.gpus.split(','))) args.num_gpus = len(gpu_ids) os.makedirs(args.save_path, exist_ok=True) print(f"Total prompts: {len(all_samples)}") print(f"GPUs: {gpu_ids}") print(f"Output: {args.save_path}") print(f"Config: {args.config_file}") # Notify user of log file location run_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") log_run_dir = os.path.join(args.save_path, "log", run_timestamp) print(f"GPU logs will be saved to: {log_run_dir}/") print(f" Each GPU will have its own log file: gpu_.log") print("") processes = [] for rank, gpu_id in enumerate(gpu_ids): print(f"[Main] Starting worker for GPU {gpu_id} (rank {rank})...") p = mp.Process(target=worker_process, args=(gpu_id, rank, args, all_samples, run_timestamp)) p.start() processes.append(p) for p in processes: p.join() if __name__ == "__main__": main()