diff --git a/apps/dreamverse/dreamverse/gpu_pool.py b/apps/dreamverse/dreamverse/gpu_pool.py index 5ecef8317..2b206c0c6 100644 --- a/apps/dreamverse/dreamverse/gpu_pool.py +++ b/apps/dreamverse/dreamverse/gpu_pool.py @@ -598,13 +598,11 @@ class GPUSlot: loop = asyncio.get_event_loop() future = loop.create_future() - # Register this request's pending future self._pending_futures[cmd.user_id] = future # Send command await loop.run_in_executor(None, self.command_queue.put, cmd) - # Ensure response reader is running await self._ensure_response_reader() # Wait for the response @@ -916,7 +914,6 @@ class GPUPool: print(f"GPU pool: {gpu_id} ready " f"({sum(1 for s in self.slots.values() if s.ready)}/{len(self.gpu_ids)})") - # Check if anyone is waiting for a GPU async with self._pool_lock: slot = self.slots[gpu_id] if slot.is_available and self.waiting_list: @@ -984,7 +981,6 @@ class GPUPool: except Exception as e: print(f"[GPU {gpu_id}] Leave user failed: {e}") - # Assign to next waiting client if GPU has capacity if slot.is_available and self.waiting_list: waiting_client_id, ready_event, _ = self.waiting_list.pop(0) self.client_gpu_map[waiting_client_id] = gpu_id diff --git a/apps/dreamverse/dreamverse/ltx2_generation.py b/apps/dreamverse/dreamverse/ltx2_generation.py index a95165e64..ed0a9e8c7 100644 --- a/apps/dreamverse/dreamverse/ltx2_generation.py +++ b/apps/dreamverse/dreamverse/ltx2_generation.py @@ -491,11 +491,9 @@ class LTX2GenerationBackend: audio_lps = (DEFAULT_LTX2_AUDIO_SAMPLE_RATE / DEFAULT_LTX2_AUDIO_HOP_LENGTH / DEFAULT_LTX2_AUDIO_DOWNSAMPLE) - # Phase 1: seed kwargs with prior-segment conditioning. self.continuation.apply_video(request_kwargs, segment_idx) self.continuation.apply_audio(request_kwargs, segment_idx, audio_lps) - # Phase 2: generate. t0 = time.perf_counter() result = self.generator.generate_video(**request_kwargs) torch.cuda.synchronize() @@ -516,7 +514,6 @@ class LTX2GenerationBackend: timings["generation_time_ms"] = result.get("generation_time", 0.0) * 1000 - # Phase 3: snapshot continuation state for the next segment. t_save_start = time.perf_counter() self.continuation.clear() self.continuation.save_video(frames) diff --git a/apps/dreamverse/dreamverse/main.py b/apps/dreamverse/dreamverse/main.py index 6768e9934..816ff888e 100644 --- a/apps/dreamverse/dreamverse/main.py +++ b/apps/dreamverse/dreamverse/main.py @@ -60,11 +60,9 @@ async def lifespan(app: FastAPI): """Application lifespan manager.""" print("Starting server...") - # Get available GPUs gpu_ids = get_available_gpus() print(f"Selected GPU ids: {gpu_ids}") - # Initialize GPU pool (spawns subprocess per GPU) runtime.gpu_pool = GPUPool(gpu_ids) await runtime.gpu_pool.initialize() diff --git a/apps/dreamverse/dreamverse/prompt_enhancer.py b/apps/dreamverse/dreamverse/prompt_enhancer.py index ab87c29c9..45dc7cf2a 100644 --- a/apps/dreamverse/dreamverse/prompt_enhancer.py +++ b/apps/dreamverse/dreamverse/prompt_enhancer.py @@ -325,7 +325,6 @@ def _parse_json_response(content: str) -> dict[str, Any]: except json.JSONDecodeError: pass - # Handle fenced blocks like ```json ... ``` with optional prose around it. fence_pattern = re.compile( r"```(?:json)?\s*([\s\S]*?)```", flags=re.IGNORECASE, diff --git a/apps/dreamverse/dreamverse/session/controller.py b/apps/dreamverse/dreamverse/session/controller.py index 42aa43b84..04eaad324 100644 --- a/apps/dreamverse/dreamverse/session/controller.py +++ b/apps/dreamverse/dreamverse/session/controller.py @@ -279,7 +279,6 @@ class SessionController: # Acquire a GPU slot. gpu_id, slot = await self.gpu_pool.acquire(client_id, websocket) - # Start session timeout. timeout_task = asyncio.create_task(session_timeout()) # Join the engine on this GPU. diff --git a/apps/fastvideo_studio/job_runner.py b/apps/fastvideo_studio/job_runner.py index e3316d833..2190b8462 100644 --- a/apps/fastvideo_studio/job_runner.py +++ b/apps/fastvideo_studio/job_runner.py @@ -347,7 +347,6 @@ class JobRunner: self._mp_manager = get_mp_context().Manager() atexit.register(self._shutdown) - # Ensure directories exist os.makedirs(self.output_dir, exist_ok=True) os.makedirs(self.log_dir, exist_ok=True) @@ -704,7 +703,6 @@ class JobRunner: raise ValueError("Job is already running") if job.status == JobStatus.COMPLETED: raise ValueError("Job already completed. Delete and re-create to run again.") - # Reset state for re-run job.status = JobStatus.PENDING job.error = None job.output_path = None @@ -841,7 +839,6 @@ class JobRunner: except Exception: logger.debug("Shutdown of the dead generator failed", exc_info=True) - # Import lazily so starting the server is fast even without a GPU. from fastvideo import VideoGenerator logger.info( @@ -1022,7 +1019,6 @@ class JobRunner: os.makedirs(os.path.join(self.output_dir, job.id), exist_ok=True) job.log_file_path = _job_log_path(self.output_dir, job.id) - # Add file handler to persist logs file_handler = logging.FileHandler(job.log_file_path, mode='w', encoding='utf-8') file_handler.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s")) job.log_file_handler = file_handler @@ -1043,7 +1039,6 @@ class JobRunner: respect_handler_level=True) queue_listener.start() - # Set output directory, create if it doesn't exist job_output_dir = os.path.join(self.output_dir, job.id) os.makedirs(job_output_dir, exist_ok=True) @@ -1167,7 +1162,6 @@ class JobRunner: finally: queue_listener.stop() - # Remove handlers and close file fastvideo_logger.removeHandler(buffer_handler) fastvideo_logger.removeHandler(file_handler) file_handler.flush() diff --git a/apps/infinite_livestream/infinite_livestream/tests/test_hls_metadata.py b/apps/infinite_livestream/infinite_livestream/tests/test_hls_metadata.py index 23aa102bd..1490d3c8f 100644 --- a/apps/infinite_livestream/infinite_livestream/tests/test_hls_metadata.py +++ b/apps/infinite_livestream/infinite_livestream/tests/test_hls_metadata.py @@ -17,7 +17,6 @@ def descriptor(name): def read_record(raw): - # Parse the TXXX UTF-8 value independently of the production serializer. assert raw[:3] == b"ID3" and raw[10:14] == b"TXXX" and raw[20] == 3 description, value = raw[21:].split(b"\x00", 1) assert description == b"infinite-livestream" diff --git a/comfyui/video_generator/video_generator.py b/comfyui/video_generator/video_generator.py index dbcdac670..d97f01d2c 100644 --- a/comfyui/video_generator/video_generator.py +++ b/comfyui/video_generator/video_generator.py @@ -150,7 +150,6 @@ class VideoGenerator: except Exception as e: print(f"Error sending interrupt signal: {e}") - # Set the interrupt event to notify other threads self._interrupt_event.set() break time.sleep(0.5) @@ -203,17 +202,14 @@ class VideoGenerator: ): print('Running FastVideo inference') - # Reset interruption flag and event self._generation_interrupted = False self._interrupt_event.clear() self._generation_result = None self._generation_exception = None - # Load pipeline config from model path pipeline_config = PipelineConfig.from_pretrained(model_path) print('pipeline_config', pipeline_config) - # Update configs with provided config dictionaries if dit_config is not None: update_config_from_args(pipeline_config.dit_config, dit_config) @@ -223,7 +219,6 @@ class VideoGenerator: if text_encoder_config is not None: update_config_from_args(pipeline_config.text_encoder_configs, text_encoder_config) - # Update top-level pipeline config with remaining arguments raw_pipeline_args = {} if embedded_cfg_scale is not None: raw_pipeline_args['embedded_cfg_scale'] = embedded_cfg_scale @@ -264,13 +259,11 @@ class VideoGenerator: print('inference_args', inference_args) - # Start a thread to run the generation self._generation_thread = threading.Thread(target=self._run_generation, args=(prompt, output_path, inference_args), daemon=True) self._generation_thread.start() - # Start a background thread to monitor for interruptions self._generation_active = True self._interrupt_thread = threading.Thread(target=self._monitor_for_interruption, daemon=True) self._interrupt_thread.start() diff --git a/examples/inference/basic/basic.py b/examples/inference/basic/basic.py index 90a702041..2045cc325 100644 --- a/examples/inference/basic/basic.py +++ b/examples/inference/basic/basic.py @@ -1,6 +1,5 @@ from fastvideo import VideoGenerator -# from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples" @@ -19,7 +18,6 @@ def main(): vae_cpu_offload=False, text_encoder_cpu_offload=True, pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" - # image_encoder_cpu_offload=False, ) # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") @@ -30,10 +28,7 @@ def main(): "wide with interest. The playful yet serene atmosphere is complemented by soft " "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True) - # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/") - # Generate another video with a different prompt, without reloading the - # model! prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " "glistening under the warm afternoon sun. The tall grass ripples gently in " "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " diff --git a/examples/inference/basic/basic_cosmos2_5_t2w.py b/examples/inference/basic/basic_cosmos2_5_t2w.py index 13590dfdf..e03220b21 100644 --- a/examples/inference/basic/basic_cosmos2_5_t2w.py +++ b/examples/inference/basic/basic_cosmos2_5_t2w.py @@ -16,7 +16,6 @@ def main(): pin_cpu_memory=True, ) - # Load default sampling parameters (negative_prompt, resolution, steps, etc.) sampling_param = SamplingParam.from_pretrained(model_path) prompt = ( diff --git a/examples/inference/basic/basic_dmd.py b/examples/inference/basic/basic_dmd.py index b0340e1c1..0612c2052 100644 --- a/examples/inference/basic/basic_dmd.py +++ b/examples/inference/basic/basic_dmd.py @@ -38,8 +38,6 @@ def main(): end_time = time.perf_counter() gen_time = end_time - start_time - # Generate another video with a different prompt, without reloading the - # model! prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " "glistening under the warm afternoon sun. The tall grass ripples gently in " "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " diff --git a/examples/inference/basic/basic_hy15.py b/examples/inference/basic/basic_hy15.py index a9dc3d141..7ee8d5d07 100644 --- a/examples/inference/basic/basic_hy15.py +++ b/examples/inference/basic/basic_hy15.py @@ -1,6 +1,5 @@ from fastvideo import VideoGenerator import json -# from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples_hy15" @@ -19,7 +18,6 @@ def main(): vae_cpu_offload=True, text_encoder_cpu_offload=True, pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" - # image_encoder_cpu_offload=False, ) prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " diff --git a/examples/inference/basic/basic_hy15_1080p.py b/examples/inference/basic/basic_hy15_1080p.py index 29ce70152..e7824addc 100644 --- a/examples/inference/basic/basic_hy15_1080p.py +++ b/examples/inference/basic/basic_hy15_1080p.py @@ -1,6 +1,5 @@ from fastvideo import VideoGenerator import json -# from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples_hy15_1080p" @@ -20,7 +19,6 @@ def main(): vae_cpu_offload=True, text_encoder_cpu_offload=True, pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" - # image_encoder_cpu_offload=False, ) prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " diff --git a/examples/inference/basic/basic_hyworld.py b/examples/inference/basic/basic_hyworld.py index 72bcf6593..4813dbad9 100644 --- a/examples/inference/basic/basic_hyworld.py +++ b/examples/inference/basic/basic_hyworld.py @@ -31,7 +31,6 @@ def main(): print(f"Num frames: {args.num_frames}") print(f"Output path: {args.output_path}") - # Initialize generator print("\nInitializing VideoGenerator for HYWorld...") generator = VideoGenerator.from_pretrained( "FastVideo/HY-WorldPlay-Bidirectional-Diffusers", @@ -44,7 +43,6 @@ def main(): image_encoder_cpu_offload=True, ) - # Generate video # The pose string is automatically converted to camera matrices by the pipeline print("\nGenerating video...") generator.generate_video( diff --git a/examples/inference/basic/basic_kandinsky5_i2v.py b/examples/inference/basic/basic_kandinsky5_i2v.py index a11dde536..7db0595cb 100644 --- a/examples/inference/basic/basic_kandinsky5_i2v.py +++ b/examples/inference/basic/basic_kandinsky5_i2v.py @@ -16,7 +16,6 @@ def main(): vae_cpu_offload=False, text_encoder_cpu_offload=True, pin_cpu_memory=True, - # image_encoder_cpu_offload=False, ) prompt = ("A woman stands up and walks away") diff --git a/examples/inference/basic/basic_kandinsky5_t2v.py b/examples/inference/basic/basic_kandinsky5_t2v.py index 37ee149a0..d6aafaafa 100644 --- a/examples/inference/basic/basic_kandinsky5_t2v.py +++ b/examples/inference/basic/basic_kandinsky5_t2v.py @@ -15,7 +15,6 @@ def main(): vae_cpu_offload=False, text_encoder_cpu_offload=True, pin_cpu_memory=True, - # image_encoder_cpu_offload=False, ) prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " diff --git a/examples/inference/basic/basic_lingbotworld_base_cam.py b/examples/inference/basic/basic_lingbotworld_base_cam.py index 2be36e022..f6490a294 100644 --- a/examples/inference/basic/basic_lingbotworld_base_cam.py +++ b/examples/inference/basic/basic_lingbotworld_base_cam.py @@ -1,7 +1,6 @@ from fastvideo import VideoGenerator from fastvideo.models.dits.lingbotworld.cam_utils import prepare_camera_embedding -# from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples_lingbotworld" @@ -18,9 +17,7 @@ def main(): dit_cpu_offload=True, # DiT need to be offloaded for MoE vae_cpu_offload=False, text_encoder_cpu_offload=True, - # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer pin_cpu_memory=True, - # image_encoder_cpu_offload=False, ) num_frames = 81 diff --git a/examples/inference/basic/basic_ltx2_distilled_fast_profile.py b/examples/inference/basic/basic_ltx2_distilled_fast_profile.py index b4e6c42b2..45640620c 100644 --- a/examples/inference/basic/basic_ltx2_distilled_fast_profile.py +++ b/examples/inference/basic/basic_ltx2_distilled_fast_profile.py @@ -258,8 +258,6 @@ def main() -> None: width=benchmark_entry.get("width", 1920), num_frames=121, num_inference_steps=5, - # image_path="examples/inference/basic/prompt1.png", - # ltx2_image_crf=0.0 ) if os.environ.get("FASTVIDEO_STAGE_LOGGING") == "0": torch.cuda.synchronize() diff --git a/examples/inference/basic/basic_matrixgame2.py b/examples/inference/basic/basic_matrixgame2.py index 3ffd39011..b6c551f7a 100644 --- a/examples/inference/basic/basic_matrixgame2.py +++ b/examples/inference/basic/basic_matrixgame2.py @@ -56,9 +56,7 @@ def main(): dit_cpu_offload=True, # DiT need to be offloaded for MoE vae_cpu_offload=False, text_encoder_cpu_offload=True, - # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer pin_cpu_memory=True, - # image_encoder_cpu_offload=False, ) num_frames = 597 diff --git a/examples/inference/basic/basic_matrixgame2_streaming.py b/examples/inference/basic/basic_matrixgame2_streaming.py index 454dafb29..d3e008988 100644 --- a/examples/inference/basic/basic_matrixgame2_streaming.py +++ b/examples/inference/basic/basic_matrixgame2_streaming.py @@ -65,7 +65,6 @@ async def main(): text_encoder_cpu_offload=True, # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer pin_cpu_memory=True, - # image_encoder_cpu_offload=False, ) max_blocks = 50 @@ -99,7 +98,6 @@ async def main(): if (await asyncio.to_thread(input, "\nContinue? (y/n): ")).lower() == 'n': break - # Save final video generator.finalize() generator.shutdown() diff --git a/examples/inference/basic/basic_mps.py b/examples/inference/basic/basic_mps.py index f1d9c214d..bebe0eb94 100644 --- a/examples/inference/basic/basic_mps.py +++ b/examples/inference/basic/basic_mps.py @@ -17,7 +17,6 @@ def main(): num_gpus=1, ) - # Create sampling parameters with reduced number of frames sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") sampling_param.num_frames = 25 # Reduce from default 81 to 25 frames bc we have to use the SDPA attn backend for mps sampling_param.height = 256 diff --git a/examples/inference/basic/basic_ray.py b/examples/inference/basic/basic_ray.py index 6d72b430c..238e93be9 100644 --- a/examples/inference/basic/basic_ray.py +++ b/examples/inference/basic/basic_ray.py @@ -1,6 +1,5 @@ from fastvideo import VideoGenerator -# from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples" @@ -20,17 +19,13 @@ def main(): text_encoder_cpu_offload=True, pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" distributed_executor_backend="ray", - # image_encoder_cpu_offload=False, ) - # Generate videos with the same simple API, regardless of GPU count prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " "wide with interest. The playful yet serene atmosphere is complemented by soft " "natural light filtering through the petals. Mid-shot, warm and cheerful tones.") video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True) - # Generate another video with a different prompt, without reloading the - # model! prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " "glistening under the warm afternoon sun. The tall grass ripples gently in " "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " diff --git a/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py b/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py index e75373b54..ec19014f9 100644 --- a/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py +++ b/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py @@ -2,7 +2,6 @@ from fastvideo import VideoGenerator, SamplingParam import json -# from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples_self_forcing_causal_wan2_2_14B_i2v" @@ -22,9 +21,7 @@ def main(): vae_cpu_offload=False, text_encoder_cpu_offload=True, dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125], - # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer pin_cpu_memory=True, - # image_encoder_cpu_offload=False, ) sampling_param = SamplingParam.from_pretrained("FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers") diff --git a/examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py b/examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py index 508f84ec2..243cd950a 100644 --- a/examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py +++ b/examples/inference/basic/basic_self_forcing_causal_wan2_2_t2v.py @@ -2,7 +2,6 @@ from fastvideo import VideoGenerator -# from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples_self_forcing_causal_wan2_2_14B_t2v" @@ -28,7 +27,6 @@ def main(): init_weights_from_safetensors_2= "/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_2_inference_transformer/", num_frame_per_block=7, - # image_encoder_cpu_offload=False, ) # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") diff --git a/examples/inference/basic/basic_wan2_2.py b/examples/inference/basic/basic_wan2_2.py index 02d8d8d38..8e115de9b 100644 --- a/examples/inference/basic/basic_wan2_2.py +++ b/examples/inference/basic/basic_wan2_2.py @@ -1,6 +1,5 @@ from fastvideo import VideoGenerator -# from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples_wan2_2_14B_t2v" @@ -20,7 +19,6 @@ def main(): text_encoder_cpu_offload=True, # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer pin_cpu_memory=True, - # image_encoder_cpu_offload=False, ) # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") @@ -36,10 +34,7 @@ def main(): height=720, width=1280, num_frames=81) - # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/") - # Generate another video with a different prompt, without reloading the - # model! prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " "glistening under the warm afternoon sun. The tall grass ripples gently in " "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " diff --git a/examples/inference/basic/basic_wan2_2_Fun.py b/examples/inference/basic/basic_wan2_2_Fun.py index c41699a46..a4c58ecc2 100644 --- a/examples/inference/basic/basic_wan2_2_Fun.py +++ b/examples/inference/basic/basic_wan2_2_Fun.py @@ -1,6 +1,5 @@ from fastvideo import VideoGenerator -# from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples_wan2_1_Fun" OUTPUT_NAME = "wan2.1_test" @@ -22,13 +21,10 @@ def main(): text_encoder_cpu_offload=True, # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer pin_cpu_memory=True, - # image_encoder_cpu_offload=False, ) prompt = "一位年轻女性穿着一件粉色的连衣裙，裙子上有白色的装饰和粉色的纽扣。她的头发是紫色的，头上戴着一个红色的大蝴蝶结，显得非常可爱和精致。她还戴着一个红色的领结，整体造型充满了少女感和活力。她的表情温柔，双手轻轻交叉放在身前，姿态优雅。背景是简单的灰色，没有任何多余的装饰，使得人物更加突出。她的妆容清淡自然，突显了她的清新气质。整体画面给人一种甜美、梦幻的感觉，仿佛置身于童话世界中。" negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" - # prompt = "A young woman with beautiful, clear eyes and blonde hair stands in the forest, wearing a white dress and a crown. Her expression is serene, reminiscent of a movie star, with fair and youthful skin. Her brown long hair flows in the wind. The video quality is very high, with a clear view. High quality, masterpiece, best quality, high resolution, ultra-fine, fantastical." - # negative_prompt = "Twisted body, limb deformities, text captions, comic, static, ugly, error, messy code." image_path = "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/wan_fun/asset_Wan2_2/v1.0/8.png" control_video_path = "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/wan_fun/asset_Wan2_2/v1.0/pose.mp4" diff --git a/examples/inference/basic/basic_wan2_2_i2v.py b/examples/inference/basic/basic_wan2_2_i2v.py index a67d57b84..d1dc6fe14 100644 --- a/examples/inference/basic/basic_wan2_2_i2v.py +++ b/examples/inference/basic/basic_wan2_2_i2v.py @@ -1,6 +1,5 @@ from fastvideo import VideoGenerator -# from fastvideo.api.sampling_param import SamplingParam OUTPUT_PATH = "video_samples_wan2_2_14B_i2v" @@ -18,9 +17,7 @@ def main(): dit_cpu_offload=True, # DiT need to be offloaded for MoE vae_cpu_offload=False, text_encoder_cpu_offload=True, - # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer pin_cpu_memory=True, - # image_encoder_cpu_offload=False, ) prompt = "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside." diff --git a/examples/inference/basic/basic_wan2_2_ti2v.py b/examples/inference/basic/basic_wan2_2_ti2v.py index c9ca208e4..311fc81d4 100644 --- a/examples/inference/basic/basic_wan2_2_ti2v.py +++ b/examples/inference/basic/basic_wan2_2_ti2v.py @@ -18,7 +18,6 @@ def main(): vae_cpu_offload=False, text_encoder_cpu_offload=True, pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" - # image_encoder_cpu_offload=False, ) # I2V is triggered just by passing in an image_path argument @@ -26,8 +25,6 @@ def main(): image_path = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG" video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, image_path=image_path) - # Generate another video with a different prompt, without reloading the - # model! # T2V mode prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame " diff --git a/examples/inference/eval/all_metrics_demo.py b/examples/inference/eval/all_metrics_demo.py index 0b6eda9c2..77cf66f1b 100644 --- a/examples/inference/eval/all_metrics_demo.py +++ b/examples/inference/eval/all_metrics_demo.py @@ -31,9 +31,6 @@ from pathlib import Path import torch -# --------------------------------------------------------------------- -# Generation — make one LTX2 video to evaluate. -# --------------------------------------------------------------------- PROMPT = ("A warm sunny backyard. The camera starts in a tight cinematic close-up " "of a woman and a man in their 30s, facing each other with serious " @@ -71,9 +68,6 @@ def generate_one_ltx2_video() -> str: return OUTPUT_PATH -# --------------------------------------------------------------------- -# Eval — the point of the script. 4 lines from "two paths" to results. -# --------------------------------------------------------------------- def _all_registered_metrics() -> list[str]: @@ -110,7 +104,6 @@ def score_all_metrics(video_path: str) -> None: results = ev.evaluate(samples=samples) torch.cuda.synchronize() t_eval1 = time.perf_counter() - # =================================== peak_gb = torch.cuda.max_memory_allocated() / 1024**3 peak_reserved_gb = torch.cuda.max_memory_reserved() / 1024**3 diff --git a/examples/inference/gradio/local/gradio_local_demo_matrixgame2.py b/examples/inference/gradio/local/gradio_local_demo_matrixgame2.py index f39c1b81b..294932c72 100644 --- a/examples/inference/gradio/local/gradio_local_demo_matrixgame2.py +++ b/examples/inference/gradio/local/gradio_local_demo_matrixgame2.py @@ -85,10 +85,6 @@ CAMERA_MAP_GTA = { def setup_model_environment(model_path: str) -> None: - # if "fullattn" in model_path.lower(): - # os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" - # else: - # os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN" os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" os.environ["FASTVIDEO_STAGE_LOGGING"] = "1" @@ -444,7 +440,6 @@ def create_gradio_interface(generators: dict[str, StreamingVideoGenerator], load # If already initialized, clean up first if state.get("initialized"): try: - # Clear accumulated frames without saving generator.accumulated_frames = [] generator.executor.execute_streaming_clear() except Exception as e: @@ -496,30 +491,22 @@ def create_gradio_interface(generators: dict[str, StreamingVideoGenerator], load if not state.get("initialized"): return state, state.get("seed", 0), "Block: 0 / 50", None, gr.update(), gr.update() - # total_start_time = time.time() config = VARIANT_CONFIG.get(model_name) generator = generators.get(config["model_path"]) mode = state["mode"] frames_per_block = state["frames_per_block"] - # Parse inputs to tensors action = get_action_tensors(mode, keyboard_key, mouse_key) keyboard_cond, mouse_cond = expand_action_to_frames(action, frames_per_block) - # run step async - # inference_start_time = time.time() frames, block_future = await generator.step_async(keyboard_cond, mouse_cond) - # inference_time = time.time() - inference_start_time # wait for block file to be written block_path = await asyncio.to_thread(block_future.result) if block_future else None state["block_idx"] = generator.block_idx block_str = f"Block: {state['block_idx']} / {state['max_blocks']}" - # total_time = time.time() - total_start_time - # Timing breakdown - # timing_html = create_timing_display(inference_time, total_time, [], frames_per_block) return state, state.get("seed", 0), block_str, block_path, gr.update(), gr.update() @@ -573,7 +560,6 @@ def main(): parser.add_argument("--port", type=int, default=7860) args = parser.parse_args() - # Load the selected model config = VARIANT_CONFIG[args.model] model_path = config["model_path"] diff --git a/examples/inference/gradio/serving/gradio_frontend.py b/examples/inference/gradio/serving/gradio_frontend.py index 23f78f853..292aaa1b5 100644 --- a/examples/inference/gradio/serving/gradio_frontend.py +++ b/examples/inference/gradio/serving/gradio_frontend.py @@ -178,7 +178,6 @@ def load_example_prompts(): print(f"Warning: Could not read {filepath}: {e}") return prompts, labels - # Load prompts from prompts.txt examples, example_labels = load_from_file("examples/inference/gradio/serving/prompts.txt") if not examples: @@ -187,7 +186,6 @@ def load_example_prompts(): ] example_labels = ["Crowded rooftop bar at night"] - # Load image mappings from JSON file prompt_to_image = {} # Try to find the JSON file relative to project root possible_json_paths = [ @@ -217,7 +215,6 @@ def load_example_prompts(): except Exception as e: print(f"Warning: Could not load image mappings from {json_path}: {e}") - # Create image paths list matching the prompts example_images = [] for prompt in examples: # Try exact match first @@ -245,17 +242,14 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling def generate_video(prompt, negative_prompt, use_negative_prompt, guidance_scale, num_frames, height, width, model_selection, input_image, progress): - # Use default seed value (randomize_seed disabled) seed = 1000 randomize_seed = False if not client.check_health(): return None, f"Backend is not available. Please check if Ray Serve is running at {backend_url}", "" - # Check if I2V model requires an image if is_i2v_model(model_selection) and not input_image: return None, "I2V models require an input image. Please upload an image.", "" - # Validate dimensions max_pixels = 720 * 1280 if height * width > max_pixels: return None, f"Video dimensions too large. Maximum: 720x1280 pixels", "" @@ -298,11 +292,6 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling if response.get("success", False): video_data = response.get("video_data", "") used_seed = response.get("seed", seed) - # inference_time = response.get("inference_time", 0.0) - # encoding_time = response.get("encoding_time", 0.0) - # total_time = response.get("total_time", 0.0) - # network_time = response.get("network_time", 0.0) - # stage_execution_times = response.get("stage_execution_times", []) # timing_details = create_timing_display( # inference_time, encoding_time, network_time, total_time, @@ -339,15 +328,6 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling ) def get_default_values(model_name): - # model_path = MODEL_PATH_MAPPING.get(model_name) - # if model_path and model_path in default_params: - # params = default_params[model_path] - # return { - # 'height': params.height, - # 'width': params.width, - # 'num_frames': params.num_frames, - # 'guidance_scale': params.guidance_scale, - # } return { 'height': 480, @@ -355,13 +335,11 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling 'num_frames': 73, } - # Get available models based on what's loaded available_models = [] for model_name, model_path in MODEL_PATH_MAPPING.items(): if model_path in default_params: available_models.append(model_name) - # Select first available model as default default_model = available_models[0] if available_models else "FastWan2.1-T2V-1.3B" initial_values = get_default_values(default_model) initial_show_image = is_i2v_model(default_model) @@ -414,7 +392,6 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling with gr.Row(): with gr.Column(): error_output = gr.Text(label="Error", visible=False) - # timing_display = gr.Markdown(label="Timing Breakdown", visible=False) with gr.Row(equal_height=False): with gr.Column(scale=1): @@ -459,7 +436,6 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling visible=False, ) - # randomize_seed = gr.Checkbox(label="Randomize seed", value=False) seed_output = gr.Number(label="Used Seed", value=1000) with gr.Column(scale=1): @@ -612,12 +588,10 @@ def main(): default_params = {} - # Load T2V model params t2v_paths = [p.strip() for p in args.t2v_model_paths.split(",") if p.strip()] for model_path in t2v_paths: default_params[model_path] = SamplingParam.from_pretrained(model_path) - # Load I2V model params i2v_paths = [p.strip() for p in args.i2v_model_paths.split(",") if p.strip()] for model_path in i2v_paths: default_params[model_path] = SamplingParam.from_pretrained(model_path) diff --git a/examples/inference/gradio/serving/ray_serve_backend.py b/examples/inference/gradio/serving/ray_serve_backend.py index 4be1518b4..66a8b1c15 100644 --- a/examples/inference/gradio/serving/ray_serve_backend.py +++ b/examples/inference/gradio/serving/ray_serve_backend.py @@ -106,13 +106,11 @@ def save_image_from_base64(image_data: str, output_dir: str) -> Optional[str]: return None try: - # Remove data URL prefix if present if image_data.startswith('data:image/'): image_data = image_data.split(',')[1] image_bytes = base64.b64decode(image_data) - # Save to temporary file os.makedirs(output_dir, exist_ok=True) temp_image_path = os.path.join(output_dir, f"temp_input_{int(time.time() * 1000)}.png") @@ -127,10 +125,6 @@ def save_image_from_base64(image_data: str, output_dir: str) -> Optional[str]: def setup_model_environment(model_path: str) -> None: - # if "fullattn" in model_path.lower(): - # os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" - # else: - # os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN" os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" os.environ["FASTVIDEO_STAGE_LOGGING"] = "1" @@ -211,7 +205,6 @@ class BaseModelDeployment: params = prepare_sampling_params(video_request, self.default_params) - # Save image if provided (for I2V) image_path = None if video_request.image_data: image_path = save_image_from_base64(video_request.image_data, self.output_path) @@ -241,7 +234,6 @@ class BaseModelDeployment: total_time = time.time() - total_start_time - # Clean up temporary image file if image_path and os.path.exists(image_path): try: os.remove(image_path) @@ -309,7 +301,6 @@ class FastVideoAPI: self.i2v_deployments = i2v_deployments or {} self.all_deployments = {**self.t2v_deployments, **self.i2v_deployments} - # Initialize Prometheus metrics self.request_count = Counter('fastvideo_requests_total', 'Total FastVideo requests', ['model_type', 'status']) self.request_duration = Histogram('fastvideo_request_duration_seconds', 'FastVideo request duration', ['model_type']) @@ -394,26 +385,21 @@ def start_ray_serve( if not ray.is_initialized(): ray.init() - # Parse T2V models t2v_paths = [p.strip() for p in t2v_model_paths.split(",") if p.strip()] t2v_reps = [int(r.strip()) for r in t2v_model_replicas.split(",") if r.strip()] if t2v_model_replicas else [] - # Parse I2V models i2v_paths = [p.strip() for p in i2v_model_paths.split(",") if p.strip()] i2v_reps = [int(r.strip()) for r in i2v_model_replicas.split(",") if r.strip()] if i2v_model_replicas else [] - # Validate configurations all_paths = t2v_paths + i2v_paths all_replicas = t2v_reps + i2v_reps validate_configuration(all_paths, all_replicas) - # Create T2V deployments t2v_deps = {} for model_path, replica_count in zip(t2v_paths, t2v_reps): t2v_dep = T2VModelDeployment.options(num_replicas=replica_count).bind(model_path, output_path) t2v_deps[model_path] = t2v_dep - # Create I2V deployments i2v_deps = {} for model_path, replica_count in zip(i2v_paths, i2v_reps): i2v_dep = I2VModelDeployment.options(num_replicas=replica_count).bind(model_path, output_path) @@ -462,7 +448,6 @@ if __name__ == "__main__": args = parser.parse_args() - # Parse and validate all models t2v_paths = [p.strip() for p in args.t2v_model_paths.split(",") if p.strip()] t2v_reps = [int(r.strip()) for r in args.t2v_model_replicas.split(",") if r.strip()] if args.t2v_model_replicas else [] diff --git a/examples/inference/gradio/serving/start_ray_serve_app.py b/examples/inference/gradio/serving/start_ray_serve_app.py index ea799db48..70e62a262 100644 --- a/examples/inference/gradio/serving/start_ray_serve_app.py +++ b/examples/inference/gradio/serving/start_ray_serve_app.py @@ -100,7 +100,6 @@ class ServiceManager: "port": self.args.backend_port } - # Add I2V parameters if provided if self.args.i2v_model_paths: backend_args["i2v_model_paths"] = self.args.i2v_model_paths if self.args.i2v_model_replicas: @@ -117,7 +116,6 @@ class ServiceManager: "port": self.args.frontend_port } - # Add I2V parameters if provided if self.args.i2v_model_paths: frontend_args["i2v_model_paths"] = self.args.i2v_model_paths diff --git a/examples/inference/lora/wan_lora_inference.py b/examples/inference/lora/wan_lora_inference.py index 1b42be28b..d47f65d64 100644 --- a/examples/inference/lora/wan_lora_inference.py +++ b/examples/inference/lora/wan_lora_inference.py @@ -5,7 +5,6 @@ OUTPUT_PATH = "./lora_out" def main(): - # Initialize VideoGenerator with the Wan model generator = VideoGenerator.from_pretrained( "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", num_gpus=1, @@ -23,13 +22,11 @@ def main(): "num_inference_steps": 32, "seed": 42, } - # Generate video with LoRA style prompt = "steamboat willie style, golden era animation, close-up of a short fluffy monster kneeling beside a melting red candle. the mood is one of wonder and curiosity, as the monster gazes at the flame with wide eyes and open mouth. Its pose and expression convey a sense of innocence and playfulness, as if it is exploring the world around it for the first time. The use of warm colors and dramatic lighting further enhances the cozy atmosphere of the image." negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" video = generator.generate_video( prompt, - # sampling_param=sampling_param, output_path=OUTPUT_PATH, save_video=True, negative_prompt=negative_prompt, diff --git a/examples/inference/lora/wan_lora_inference_from_ckpt.py b/examples/inference/lora/wan_lora_inference_from_ckpt.py index 99f6bb1c3..31a488887 100644 --- a/examples/inference/lora/wan_lora_inference_from_ckpt.py +++ b/examples/inference/lora/wan_lora_inference_from_ckpt.py @@ -8,7 +8,6 @@ OUTPUT_PATH = "./lora_out" def main(): - # Initialize VideoGenerator with the Wan model generator = VideoGenerator.from_pretrained( "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", num_gpus=1, @@ -27,7 +26,6 @@ def main(): "num_inference_steps": 50, "seed": 42, } - # Generate video with LoRA style prompt = "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press." video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, **kwargs) diff --git a/examples/inference/optimizations/attention_example.py b/examples/inference/optimizations/attention_example.py index ba8e830a5..7872874a8 100644 --- a/examples/inference/optimizations/attention_example.py +++ b/examples/inference/optimizations/attention_example.py @@ -5,7 +5,6 @@ from fastvideo import VideoGenerator def main(): - # set the attention backend os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" start_time = time.perf_counter() diff --git a/fastvideo/attention/backends/abstract.py b/fastvideo/attention/backends/abstract.py index 1d25a7c19..265a99e8f 100644 --- a/fastvideo/attention/backends/abstract.py +++ b/fastvideo/attention/backends/abstract.py @@ -50,14 +50,7 @@ class AttentionBackend(ABC): def get_metadata_cls() -> type["AttentionMetadata"]: raise NotImplementedError - # @staticmethod - # @abstractmethod - # def get_state_cls() -> Type["AttentionState"]: - # raise NotImplementedError - # @classmethod - # def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata": - # return cls.get_metadata_cls()(*args, **kwargs) @staticmethod @abstractmethod diff --git a/fastvideo/attention/backends/bsa_attn.py b/fastvideo/attention/backends/bsa_attn.py index c505d04f2..cac98ec0b 100644 --- a/fastvideo/attention/backends/bsa_attn.py +++ b/fastvideo/attention/backends/bsa_attn.py @@ -43,9 +43,6 @@ logger = init_logger(__name__) BSA_TILE_SIZE = (4, 4, 4) -# --------------------------------------------------------------------------- -# Cached index helpers (same pattern as VSA) -# --------------------------------------------------------------------------- @functools.lru_cache(maxsize=10) @@ -80,9 +77,6 @@ def get_reverse_tile_partition_indices( return torch.argsort(get_tile_partition_indices(dit_seq_shape, tile_size, device)) -# --------------------------------------------------------------------------- -# BSA core operations -# --------------------------------------------------------------------------- def _prune_queries( @@ -286,7 +280,6 @@ def _compute_sparse_attention_flash( # Per-head path: process each head individually for h in range(H): head_mask = kv_mask[b, h] # [N, N] - # Process single head: squeeze head dim, run FA, put back _flash_attn_single_head( sparse_q[b, h], k_blocks[b, h], @@ -537,9 +530,6 @@ def _reconstruct_pruned( return full_output -# --------------------------------------------------------------------------- -# FastVideo backend classes -# --------------------------------------------------------------------------- class BSAAttentionBackend(AttentionBackend): diff --git a/fastvideo/attention/backends/sage_attn.py b/fastvideo/attention/backends/sage_attn.py index 5db61f6ac..904309165 100644 --- a/fastvideo/attention/backends/sage_attn.py +++ b/fastvideo/attention/backends/sage_attn.py @@ -26,9 +26,6 @@ class SageAttentionBackend(AttentionBackend): def get_impl_cls() -> type["SageAttentionImpl"]: return SageAttentionImpl - # @staticmethod - # def get_metadata_cls() -> Type["AttentionMetadata"]: - # return FlashAttentionMetadata class SageAttentionImpl(AttentionImpl): diff --git a/fastvideo/attention/backends/sage_attn3.py b/fastvideo/attention/backends/sage_attn3.py index 5830c14eb..b31a1a5aa 100644 --- a/fastvideo/attention/backends/sage_attn3.py +++ b/fastvideo/attention/backends/sage_attn3.py @@ -34,9 +34,6 @@ class SageAttention3Backend(AttentionBackend): def get_builder_cls() -> type["AttentionMetadataBuilder"]: raise NotImplementedError - # @staticmethod - # def get_metadata_cls() -> Type["AttentionMetadata"]: - # return FlashAttentionMetadata class SageAttention3Impl(AttentionImpl): diff --git a/fastvideo/attention/backends/sdpa.py b/fastvideo/attention/backends/sdpa.py index 13ccf7fde..3811cfa60 100644 --- a/fastvideo/attention/backends/sdpa.py +++ b/fastvideo/attention/backends/sdpa.py @@ -32,9 +32,6 @@ class SDPABackend(AttentionBackend): def get_impl_cls() -> type["SDPAImpl"]: return SDPAImpl - # @staticmethod - # def get_metadata_cls() -> Type["AttentionMetadata"]: - # return FlashAttentionMetadata @dataclass diff --git a/fastvideo/attention/backends/sla.py b/fastvideo/attention/backends/sla.py index efb3cf882..c92151566 100644 --- a/fastvideo/attention/backends/sla.py +++ b/fastvideo/attention/backends/sla.py @@ -34,9 +34,6 @@ from fastvideo.logger import init_logger logger = init_logger(__name__) -# ============================================================================ -# SLA Utility functions (moved from sla_kernels/utils.py) -# ============================================================================ @triton.jit @@ -110,9 +107,6 @@ def get_block_map( return sparse_map, lut, topk -# ============================================================================ -# SLA Backend classes -# ============================================================================ class SLAAttentionBackend(AttentionBackend): @@ -286,20 +280,16 @@ class SLAAttentionImpl(AttentionImpl, nn.Module): """ original_dtype = query.dtype - # Convert from FastVideo format (B, L, H, D) to SLA format (B, H, L, D) q = query.transpose(1, 2).contiguous() k = key.transpose(1, 2).contiguous() v = value.transpose(1, 2).contiguous() - # Get topk ratio from metadata if available topk_ratio = self.topk_ratio if hasattr(attn_metadata, 'topk_ratio'): topk_ratio = attn_metadata.topk_ratio # type: ignore[union-attr] - # Compute block-sparse attention pattern sparse_map, lut, real_topk = get_block_map(q, k, topk_ratio=topk_ratio, BLKQ=self.BLKQ, BLKK=self.BLKK) - # Convert to compute dtype q = q.to(self.dtype) k = k.to(self.dtype) v = v.to(self.dtype) @@ -322,13 +312,11 @@ class SLAAttentionImpl(AttentionImpl, nn.Module): # Combine sparse and linear outputs output = (o_s + o_l).to(original_dtype) - # Convert back to FastVideo format (B, L, H, D) output = output.transpose(1, 2) return output -# Check if spas_sage_attn is available for SageSLA SAGESLA_ENABLED = True try: import spas_sage_attn._qattn as qattn @@ -482,12 +470,10 @@ class SageSLAAttentionImpl(AttentionImpl, nn.Module): """ original_dtype = query.dtype - # Convert from FastVideo format (B, L, H, D) to SLA format (B, H, L, D) q = query.transpose(1, 2).contiguous() k = key.transpose(1, 2).contiguous() v = value.transpose(1, 2).contiguous() - # Get topk ratio from metadata if available topk_ratio = self.topk_ratio if hasattr(attn_metadata, 'topk_ratio'): topk_ratio = attn_metadata.topk_ratio # type: ignore[union-attr] @@ -499,10 +485,8 @@ class SageSLAAttentionImpl(AttentionImpl, nn.Module): else: BLKQ, BLKK = 128, 64 - # Compute block-sparse attention pattern sparse_map, lut, real_topk = get_block_map(q, k, topk_ratio=topk_ratio, BLKQ=BLKQ, BLKK=BLKK) - # Convert to compute dtype q = q.to(self.dtype) k = k.to(self.dtype) v = v.to(self.dtype) @@ -555,7 +539,6 @@ class SageSLAAttentionImpl(AttentionImpl, nn.Module): # Combine sparse and linear outputs output = (o_s + o_l).to(original_dtype) - # Convert back to FastVideo format (B, L, H, D) output = output.transpose(1, 2) return output diff --git a/fastvideo/attention/backends/vmoba.py b/fastvideo/attention/backends/vmoba.py index 85519f8b3..a8e391aca 100644 --- a/fastvideo/attention/backends/vmoba.py +++ b/fastvideo/attention/backends/vmoba.py @@ -148,7 +148,6 @@ class VMOBAAttentionImpl(AttentionImpl): """ batch_size, sequence_length, num_heads, head_dim = query.shape - # select chunk type according to layer idx: loop_layer_num = attn_metadata.temporal_layer + attn_metadata.spatial_layer + attn_metadata.st_layer assert self.layer_idx is not None, "VMoBA attention requires layer_idx to be set" moba_layer = self.layer_idx - attn_metadata.first_full_layer @@ -179,7 +178,6 @@ class VMOBAAttentionImpl(AttentionImpl): key = rearrange(key, "b s ... -> (b s) ...") value = rearrange(value, "b s ... -> (b s) ...") - # current_timestep=attn_metadata.current_timestep hidden_states = moba_attn_varlen( query, key, diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py index a64c6d604..7cfa502e1 100644 --- a/fastvideo/attention/layer.py +++ b/fastvideo/attention/layer.py @@ -134,7 +134,6 @@ class DistributedAttention(nn.Module): - o (torch.Tensor): Output tensor after attention for the main sequence - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided """ - # Check input shapes assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors" # Stack QKV qkv = torch.cat([q, k, v], dim=0) # [3*batch, seq_len, num_heads, head_dim] @@ -186,7 +185,6 @@ class DistributedAttention(nn.Module): if freqs_cis is not None: cos, sin = freqs_cis qkv[:batch_size * 2] = _apply_rotary_emb(qkv[:batch_size * 2], cos, sin, is_neox_style=False) - # Apply backend-specific preprocess_qkv qkv = self.attn_impl.preprocess_qkv(qkv, ctx_attn_metadata) # Concatenate with replicated QKV if provided @@ -208,9 +206,7 @@ class DistributedAttention(nn.Module): split_idx = original_seq_len replicated_output = output[:, split_idx:] output = output[:, :split_idx] - # TODO: make this asynchronous replicated_output = sequence_model_parallel_all_gather(replicated_output.contiguous(), dim=2) - # Apply backend-specific postprocess_output output = self.attn_impl.postprocess_output(output, ctx_attn_metadata) output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len)) @@ -255,9 +251,7 @@ class DistributedAttention_VSA(DistributedAttention): - o (torch.Tensor): Output tensor after attention for the main sequence - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided """ - # Check text tokens are not supported for VSA now assert replicated_q is None and replicated_k is None and replicated_v is None, "Replicated QKV is not supported for VSA now" - # Check input shapes assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors" forward_context: ForwardContext = get_forward_context() @@ -300,7 +294,6 @@ class DistributedAttention_VSA(DistributedAttention): # Redistribute back if using sequence parallelism replicated_output = None - # Apply backend-specific postprocess_output output = self.attn_impl.postprocess_output(output, ctx_attn_metadata) output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len)) @@ -367,7 +360,6 @@ class LocalAttention(nn.Module): Returns: torch.Tensor: Output tensor after local attention """ - # Check input shapes assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors" forward_context: ForwardContext = get_forward_context() diff --git a/fastvideo/attention/selector.py b/fastvideo/attention/selector.py index 76202d342..411d51cb0 100644 --- a/fastvideo/attention/selector.py +++ b/fastvideo/attention/selector.py @@ -310,7 +310,6 @@ def _cached_get_attn_backend( if selected_backend is None and default_backend is not None: selected_backend = default_backend - # get device-specific attn_backend from fastvideo.platforms import current_platform if (selected_backend is not None and selected_backend not in supported_attention_backends): diff --git a/fastvideo/configs/models/dits/hyworld.py b/fastvideo/configs/models/dits/hyworld.py index 66d33e055..218cb0ddc 100644 --- a/fastvideo/configs/models/dits/hyworld.py +++ b/fastvideo/configs/models/dits/hyworld.py @@ -122,7 +122,6 @@ class HYWorldArchConfig(DiTArchConfig): def __post_init__(self): super().__post_init__() - # Convert HY-WorldPlay naming to FastVideo naming conventions self.num_attention_heads: int = self.heads_num self.attention_head_dim: int = self.hidden_size // self.heads_num self.num_layers: int = self.mm_double_blocks_depth @@ -145,7 +144,6 @@ class HYWorldArchConfig(DiTArchConfig): else: self.in_channels = self.in_channels * 2 + 1 # 32 * 2 + 1 = 65 - # Handle patch_size (can be list/tuple or int) if isinstance(self.patch_size, list | tuple): self.patch_size_t: int = self.patch_size[0] # assume square patch size for height and width @@ -154,7 +152,6 @@ class HYWorldArchConfig(DiTArchConfig): else: self.patch_size_t = 1 - # Convert qk_norm to string format if isinstance(self.qk_norm, bool): if self.qk_norm: self.qk_norm = "rms_norm" if self.qk_norm_type == "rms" else self.qk_norm_type diff --git a/fastvideo/configs/models/dits/ltx2.py b/fastvideo/configs/models/dits/ltx2.py index 04150fad5..e728fcaaf 100644 --- a/fastvideo/configs/models/dits/ltx2.py +++ b/fastvideo/configs/models/dits/ltx2.py @@ -68,7 +68,6 @@ class LTX2VideoArchConfig(DiTArchConfig): timestep_scale_multiplier: int = 1000 use_middle_indices_grid: bool = True - # Patchification (video-only path) patch_size: tuple[int, int, int] = (1, 1, 1) num_channels_latents: int = 128 in_channels: int | None = None diff --git a/fastvideo/configs/models/dits/wan_animate.py b/fastvideo/configs/models/dits/wan_animate.py index 451efef8b..13bbb0651 100644 --- a/fastvideo/configs/models/dits/wan_animate.py +++ b/fastvideo/configs/models/dits/wan_animate.py @@ -80,7 +80,6 @@ class WanAnimateArchConfig(WanVideoArchConfig): image_dim: int | None = 1280 added_kv_proj_dim: int | None = 5120 - # --- Animate-specific (names match config.json keys exactly) --- # in_channels = 2 * latent_channels + 4 mask channels; asserted below. latent_channels: int = 16 # None -> the model uses the LIA channel table keyed by feature-map size diff --git a/fastvideo/configs/pipelines/base.py b/fastvideo/configs/pipelines/base.py index 7faf283b4..978b1f4e4 100644 --- a/fastvideo/configs/pipelines/base.py +++ b/fastvideo/configs/pipelines/base.py @@ -201,11 +201,9 @@ class PipelineConfig: help="Comma-separated list of denoising steps (e.g., '1000,757,522')", ) - # Add VAE configuration arguments from fastvideo.configs.models.vaes.base import VAEConfig VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config") - # Add DiT configuration arguments from fastvideo.configs.models.dits.base import DiTConfig DiTConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}dit-config") @@ -324,7 +322,6 @@ class PipelineConfig: model_dicts.append(model_dict) output_dict[key] = model_dicts elif isinstance(value, tuple) and all(callable(f) for f in value): - # Skip dumping functions del_keys.append(key) for key in del_keys: diff --git a/fastvideo/configs/pipelines/gen3c.py b/fastvideo/configs/pipelines/gen3c.py index a98c3d016..ce3ed28c2 100644 --- a/fastvideo/configs/pipelines/gen3c.py +++ b/fastvideo/configs/pipelines/gen3c.py @@ -141,7 +141,6 @@ class Gen3CConfig(PipelineConfig): self._vae_latent_dim = 16 - # Validate frame buffer configuration matches DiT if hasattr(self.dit_config, 'arch_config'): arch_config = self.dit_config.arch_config if (hasattr(arch_config, 'frame_buffer_max') and arch_config.frame_buffer_max != self.frame_buffer_max): @@ -157,7 +156,6 @@ class Gen3CConfig(PipelineConfig): class Gen3CInferenceConfig(Gen3CConfig): """Configuration for GEN3C inference with optimized defaults.""" - # Use smaller batch sizes for inference batch_size: int = 1 # Enable gradient checkpointing for memory efficiency diff --git a/fastvideo/configs/utils.py b/fastvideo/configs/utils.py index 5222315f6..7f352fe5e 100644 --- a/fastvideo/configs/utils.py +++ b/fastvideo/configs/utils.py @@ -12,7 +12,6 @@ def update_config_from_args(config: Any, args_dict: dict[str, Any], prefix: str prefix: Prefix for the configuration parameters in the args_dict. If None, assumes direct attribute mapping without prefix. """ - # Handle top-level attributes (no prefix) args_not_to_remove = [ 'model_path', ] @@ -27,7 +26,6 @@ def update_config_from_args(config: Any, args_dict: dict[str, Any], prefix: str if pop_args: args_to_remove.append(key) else: - # Handle nested attributes with prefix prefix_with_dot = f"{prefix}." for key, value in args_dict.items(): if key.startswith(prefix_with_dot) and value is not None: diff --git a/fastvideo/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py b/fastvideo/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py index 1062f49b6..8d7c59cf2 100644 --- a/fastvideo/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py +++ b/fastvideo/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py @@ -48,7 +48,6 @@ def main() -> None: maybe_init_distributed_environment_and_model_parallel(tp_size=(world_size + 1) // 2, sp_size=(world_size + 1) // 2) logger.info("Initialized distributed environment with world_size=%d", world_size) - # Create DataLoader with proper settings dataset, dataloader = build_parquet_iterable_style_dataloader(args.path, args.batch_size, args.num_data_workers) logger.info("Initialized dataloader") @@ -62,7 +61,6 @@ def main() -> None: if i >= args.num_batches_per_epoch - 1: break - # Save dataloader state using distributed checkpoint checkpoint_dir = pathlib.Path(args.checkpoint_path) logger.info("Rank %d: Saving dataloader state to %s", get_world_rank(), checkpoint_dir) states = {"dataloader": dataloader} @@ -73,7 +71,6 @@ def main() -> None: logger.info("Rank %d: Saved checkpoint in %.2f seconds", get_world_rank(), end_time - begin_time) - # Make sure all processes wait for checkpoint to be saved if world_size > 1: dist.barrier() @@ -101,7 +98,6 @@ def main() -> None: if i >= args.num_batches_per_epoch * 2 - 1: break - # Verify all sums match if all(abs(a - b) < 1e-6 for a, b in zip(first_pass_sums, second_pass_sums, strict=True)): logger.info("All latent sums match between passes - resume verification successful!") else: @@ -115,11 +111,9 @@ def main() -> None: if i >= args.num_batches_per_epoch: break - # Move data to device latents = latents.to(get_local_torch_device()) embeddings = embeddings.to(get_local_torch_device()) - # Calculate actual batch size batch_size = latents.size(0) total_samples += batch_size total_batches += 1 diff --git a/fastvideo/dataset/benchmarks/benchmark_parquet_dataset_map_style.py b/fastvideo/dataset/benchmarks/benchmark_parquet_dataset_map_style.py index 9ada11e73..84978500a 100644 --- a/fastvideo/dataset/benchmarks/benchmark_parquet_dataset_map_style.py +++ b/fastvideo/dataset/benchmarks/benchmark_parquet_dataset_map_style.py @@ -51,7 +51,6 @@ def main() -> None: maybe_init_distributed_environment_and_model_parallel(tp_size=(world_size + 1) // 2, sp_size=(world_size + 1) // 2) logger.info("Initialized distributed environment with world_size=%d", world_size) - # Create DataLoader with proper settings dataset, dataloader = build_parquet_map_style_dataloader(args.path, args.batch_size, parquet_schema=pyarrow_schema_t2v, @@ -70,7 +69,6 @@ def main() -> None: if i >= args.num_batches_per_epoch - 1: break - # Save dataloader state using distributed checkpoint checkpoint_dir = pathlib.Path(args.checkpoint_path) logger.info("Rank %d: Saving dataloader state to %s", get_world_rank(), checkpoint_dir) states = {"dataloader": dataloader} @@ -81,7 +79,6 @@ def main() -> None: logger.info("Rank %d: Saved checkpoint in %.2f seconds", get_world_rank(), end_time - begin_time) - # Make sure all processes wait for checkpoint to be saved if world_size > 1: dist.barrier() @@ -119,7 +116,6 @@ def main() -> None: if i >= args.num_batches_per_epoch * 2 - 1: break - # Verify all sums match if all(abs(a - b) < 1e-6 for a, b in zip(first_pass_sums, second_pass_sums, strict=True)): logger.info("All latent sums match between passes - resume verification successful!") else: @@ -135,11 +131,9 @@ def main() -> None: if i >= args.num_batches_per_epoch: break - # Move data to device latents = latents.to(get_local_torch_device()) embeddings = embeddings.to(get_local_torch_device()) - # Calculate actual batch size batch_size = latents.size(0) total_samples += batch_size total_batches += 1 diff --git a/fastvideo/dataset/dataloader/parquet_io.py b/fastvideo/dataset/dataloader/parquet_io.py index 242abaefd..1e959efa1 100644 --- a/fastvideo/dataset/dataloader/parquet_io.py +++ b/fastvideo/dataset/dataloader/parquet_io.py @@ -153,7 +153,6 @@ class ParquetDatasetWriter: first_index = _next_chunk_index(self.out_dir) if remainder_table is not None and len(remainder_table) > 0: if write_remainder: - # Write the remainder as a final small file (worker_0) worker_dir = os.path.join(self.out_dir, "worker_0") os.makedirs(worker_dir, exist_ok=True) remainder_path = os.path.join(worker_dir, f"data_chunk_{first_index + total_chunks}.parquet") diff --git a/fastvideo/dataset/dataloader/record_schema.py b/fastvideo/dataset/dataloader/record_schema.py index d2024c428..8d80dd08f 100644 --- a/fastvideo/dataset/dataloader/record_schema.py +++ b/fastvideo/dataset/dataloader/record_schema.py @@ -23,11 +23,9 @@ def basic_t2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]: width = batch.width[idx] if batch.width is not None else 0 height = batch.height[idx] if batch.height is not None else 0 - # Get FPS - single value in PreprocessBatch fps_val = float(batch.fps[idx]) if batch.fps is not None else 0.0 # For duration, we need to calculate it or use a default since it's not in PreprocessBatch - # duration = num_frames / fps if available duration_val = 0.0 if batch.num_frames[idx] and batch.fps[idx] and batch.fps[idx] > 0: duration_val = float(batch.num_frames[idx]) / float(batch.fps[idx]) diff --git a/fastvideo/dataset/latent_datasets.py b/fastvideo/dataset/latent_datasets.py index 277784619..35669dfab 100644 --- a/fastvideo/dataset/latent_datasets.py +++ b/fastvideo/dataset/latent_datasets.py @@ -20,7 +20,6 @@ class LatentDataset(Dataset): self.json_path = json_path self.cfg_rate = cfg_rate self.seed = seed - # Create a seeded random generator for deterministic CFG self.rng = random.Random(seed) self.datase_dir_path = os.path.dirname(json_path) self.video_dir = os.path.join(self.datase_dir_path, "video") @@ -29,8 +28,6 @@ class LatentDataset(Dataset): self.prompt_attention_mask_dir = os.path.join(self.datase_dir_path, "prompt_attention_mask") with open(self.json_path) as f: self.data_anno = json.load(f) - # json.load(f) already keeps the order - # self.data_anno = sorted(self.data_anno, key=lambda x: x['latent_path']) self.num_latent_t = num_latent_t self.uncond_prompt_embed = torch.zeros(256, 4096).to(torch.float32) @@ -42,7 +39,6 @@ class LatentDataset(Dataset): latent_file = self.data_anno[idx]["latent_path"] prompt_embed_file = self.data_anno[idx]["prompt_embed_path"] prompt_attention_mask_file = self.data_anno[idx]["prompt_attention_mask"] - # load latent = torch.load( os.path.join(self.latent_dir, latent_file), map_location="cpu", @@ -75,7 +71,6 @@ def latent_collate_function(batch): # text_attn_mask: b 1 l # needs to check if the latent/prompt' size and apply padding & attn mask latents, prompt_embeds, prompt_attention_masks = zip(*batch, strict=True) - # calculate max shape max_t = max([latent.shape[1] for latent in latents]) max_h = max([latent.shape[2] for latent in latents]) max_w = max([latent.shape[3] for latent in latents]) @@ -96,7 +91,6 @@ def latent_collate_function(batch): ] # attn mask latent_attn_mask = torch.ones(len(latent_list), max_t, max_h, max_w) - # set to 0 if padding for i, latent in enumerate(latent_list): latent_attn_mask[i, latent.shape[1]:, :, :] = 0 latent_attn_mask[i, :, latent.shape[2]:, :] = 0 diff --git a/fastvideo/dataset/parquet_dataset_iterable_style.py b/fastvideo/dataset/parquet_dataset_iterable_style.py index 7c9b80b72..12520d028 100644 --- a/fastvideo/dataset/parquet_dataset_iterable_style.py +++ b/fastvideo/dataset/parquet_dataset_iterable_style.py @@ -17,7 +17,6 @@ logger = init_logger(__name__) class BatchIterator: - # TODO: Implement state_dict and load_state_dict to support resume. def __init__(self, files, batch_size, text_padding_length, keys, worker_num_samples, read_batch_size): self.files = files self.batch_size = batch_size @@ -77,19 +76,16 @@ class LatentsParquetIterStyleDataset(IterableDataset): self.text_padding_length = text_padding_length self.seed = seed self.read_batch_size = read_batch_size - # Get distributed training info self.global_rank = get_world_rank() self.world_size = get_world_size() self.sp_world_size = get_sp_world_size() self.num_sp_groups = self.world_size // self.sp_world_size num_workers = 1 if num_workers == 0 else num_workers - # Get sharding info shard_parquet_files, shard_total_samples, shard_parquet_lengths = shard_parquet_files_across_sp_groups_and_workers( self.path, self.num_sp_groups, num_workers, seed) if drop_last: self.worker_num_samples = min(shard_total_samples) // batch_size * batch_size - # Assign files to current rank's SP group ith_sp_group = self.global_rank // self.sp_world_size self.sp_group_parquet_files = shard_parquet_files[ith_sp_group::self.num_sp_groups] self.sp_group_parquet_lengths = shard_parquet_lengths[ith_sp_group::self.num_sp_groups] @@ -142,7 +138,6 @@ def shard_parquet_files_across_sp_groups_and_workers( - List of total samples per shard - List of dictionaries mapping file paths to their lengths """ - # Check if sharding plan already exists sharding_info_dir = os.path.join(path, f"sharding_info_{num_sp_groups}_sp_groups_{num_workers}_workers") # Only rank 0 handles cache checking and file scanning @@ -185,7 +180,6 @@ def shard_parquet_files_across_sp_groups_and_workers( if not parquet_files: raise ValueError("No parquet files found in %s", path) - # Calculate file lengths efficiently using a single pass logger.info("Calculating file lengths...") lengths = [] for file in tqdm.tqdm(parquet_files, desc="Reading parquet files"): @@ -199,7 +193,6 @@ def shard_parquet_files_across_sp_groups_and_workers( sorted_files = [parquet_files[i] for i in sorted_indices] sorted_lengths = [lengths[i] for i in sorted_indices] - # Create shards num_shards = num_sp_groups * num_workers shard_parquet_files = [[] for _ in range(num_shards)] shard_total_samples = [0] * num_shards @@ -218,7 +211,6 @@ def shard_parquet_files_across_sp_groups_and_workers( rng = random.Random(seed) rng.shuffle(shard) - # Save the sharding plan os.makedirs(sharding_info_dir, exist_ok=True) with open(os.path.join(sharding_info_dir, "shard_parquet_files.pkl"), "wb") as f: pickle.dump(shard_parquet_files, f) diff --git a/fastvideo/dataset/parquet_dataset_map_style.py b/fastvideo/dataset/parquet_dataset_map_style.py index eb35d91f4..b8be81d73 100644 --- a/fastvideo/dataset/parquet_dataset_map_style.py +++ b/fastvideo/dataset/parquet_dataset_map_style.py @@ -98,7 +98,6 @@ class DP_SP_BatchSampler(Sampler[list[int]]): if drop_first_row: self.dataset_size -= 1 - # shard the indices to each sp group ith_sp_group = self.global_rank // self.sp_world_size sp_group_local_indices = global_indices[ith_sp_group::self.num_sp_groups] self.sp_group_local_indices = sp_group_local_indices @@ -1033,7 +1032,6 @@ def get_parquet_files_and_length(path: str | Sequence[str] | dict[str, int], *, return file_names_sorted, lengths_sorted dataset_root = os.path.realpath(os.path.expanduser(specs[0][0])) - # Check if cached info exists cache_dir = os.path.join(dataset_root, "map_style_cache") cache_file = os.path.join(cache_dir, "file_info.pkl") @@ -1118,7 +1116,6 @@ def get_parquet_files_and_length(path: str | Sequence[str] | dict[str, int], *, # sort according to file name to ensure all rank has the same order file_names_sorted, lengths_sorted = zip(*sorted(zip(file_names, lengths, strict=True), key=lambda x: x[0]), strict=True) - # Save the cache os.makedirs(cache_dir, exist_ok=True) with open(cache_file, "wb") as f: pickle.dump((file_names_sorted, lengths_sorted), f) @@ -1351,10 +1348,8 @@ class LatentsParquetMapStyleDataset(Dataset): Returns the processed negative prompt data (latents, embeddings, masks, info). """ - # Read first row from first parquet file file_path = self.parquet_files[0] row_idx = 0 - # Read the negative prompt data row_dict = read_row_from_parquet_file( [file_path], row_idx, diff --git a/fastvideo/dataset/preprocessing_datasets.py b/fastvideo/dataset/preprocessing_datasets.py index 9e27f86b4..e0123bf37 100644 --- a/fastvideo/dataset/preprocessing_datasets.py +++ b/fastvideo/dataset/preprocessing_datasets.py @@ -130,12 +130,10 @@ class DataValidationStage(DatasetFilterStage): Returns: True if valid, False if invalid """ - # Check for caption if batch.cap is None: return False if batch.is_video: - # Validate video-specific fields if batch.duration is None or batch.fps is None: return False elif not batch.is_image: @@ -163,7 +161,6 @@ class FrameSamplingStage(DatasetFilterStage): self.speed_factor = speed_factor self.video_length_tolerance_range = video_length_tolerance_range self.drop_short_ratio = drop_short_ratio - # Create a seeded random generator for deterministic sampling self.rng = random.Random(seed) def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool: @@ -184,7 +181,6 @@ class FrameSamplingStage(DatasetFilterStage): num_frames = math.ceil(batch.fps * batch.duration) - # Check if video is too long if (num_frames / batch.fps > self.video_length_tolerance_range * (self.num_frames / self.train_fps * self.speed_factor)): return False @@ -313,7 +309,6 @@ class TextEncodingStage(DatasetStage): self.tokenizer = tokenizer self.text_max_length = text_max_length self.cfg_rate = cfg_rate - # Create a seeded random generator for deterministic CFG self.rng = random.Random(seed) def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch: @@ -397,7 +392,6 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, torch.distribu self.temporal_sample = temporal_sample self.seed = seed - # Initialize tokenizer tokenizer_path = os.path.join(args.model_path, "tokenizer") tokenizer = None if os.path.exists(tokenizer_path): @@ -406,10 +400,8 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, torch.distribu except (ValueError, OSError): pass - # Initialize processing stages self._init_stages(args, transform, transform_topcrop, tokenizer) - # Process metadata self.processed_batches = self._process_metadata() def _init_stages(self, args, transform, transform_topcrop, tokenizer) -> None: @@ -433,7 +425,6 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, torch.distribu def _load_raw_data(self) -> list[dict]: """Load raw data from JSON files.""" - # Read folder-annotation pairs with open(self.data_merge_path) as f: folder_anno_pairs = [line.strip().split(",") for line in f if line.strip()] assert len(folder_anno_pairs) == 1, "Only support one folder-annotation pair" @@ -444,7 +435,6 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, torch.distribu with open(annotation_file) as f: data_items = json.load(f) - # Update paths with folder prefix for item in data_items: item["path"] = opj(folder, item["path"]) if "action_path" in item and item["action_path"]: @@ -457,7 +447,6 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, torch.distribu raw_data = self._load_raw_data() processed_batches = [] - # Initialize counters filter_counts = {"validation_failed": 0, "frame_sampling_failed": 0} sample_num_frames: list[int] = [] @@ -469,11 +458,9 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, torch.distribu duration=item.get("duration"), action_path=item.get("action_path")) - # Apply filtering stages if not self._apply_filter_stages(batch, filter_counts): continue - # Apply frame sampling processing batch = self.frame_sampling_stage.process(batch, temporal_sample_fn=self.temporal_sample) processed_batches.append(batch) @@ -515,13 +502,11 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, torch.distribu """Get a single processed data item.""" batch = replace(self.processed_batches[idx]) - # Apply transformation stages batch = self.video_transform_stage.process(batch) batch = self.image_transform_stage.process(batch) if self.text_encoding_stage is not None: batch = self.text_encoding_stage.process(batch) - # Build result dictionary result = { "pixel_values": batch.pixel_values, "path": batch.path, @@ -532,11 +517,9 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, torch.distribu result["input_ids"] = batch.input_ids result["cond_mask"] = batch.cond_mask - # Add video-specific fields if batch.is_video: result.update({"fps": batch.fps, "duration": batch.duration}) - # Add action_path if batch.action_path: result["action_path"] = batch.action_path @@ -570,17 +553,14 @@ class TextDataset(torch.utils.data.IterableDataset, torch.distributed.checkpoint self.args = args self.seed = seed - # Initialize tokenizer tokenizer_path = os.path.join(args.model_path, "tokenizer") tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, cache_dir=args.cache_dir) - # Initialize text encoding stage self.text_encoding_stage = TextEncodingStage(tokenizer=tokenizer, text_max_length=args.text_max_length, cfg_rate=getattr(args, 'training_cfg_rate', 0.0), seed=self.seed) - # Process text data self.processed_batches = self._process_text_data() def _load_text_data(self) -> list[str]: @@ -619,7 +599,6 @@ class TextDataset(torch.utils.data.IterableDataset, torch.distributed.checkpoint def __iter__(self): """Iterator for the dataset.""" - # Set up distributed sampling if needed if torch.distributed.is_available() and torch.distributed.is_initialized(): rank = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() @@ -627,7 +606,6 @@ class TextDataset(torch.utils.data.IterableDataset, torch.distributed.checkpoint rank = 0 world_size = 1 - # Calculate chunk for this rank total_items = len(self.processed_batches) items_per_rank = math.ceil(total_items / world_size) start_idx = rank * items_per_rank + self.start_idx @@ -642,7 +620,6 @@ class TextDataset(torch.utils.data.IterableDataset, torch.distributed.checkpoint """Get a single processed text item.""" batch = self.processed_batches[idx] - # Apply text encoding stage batch = self.text_encoding_stage.process(batch) # Build result dictionary for text-only processing with required schema fields diff --git a/fastvideo/dataset/transform.py b/fastvideo/dataset/transform.py index c9a215e8d..676c459ec 100644 --- a/fastvideo/dataset/transform.py +++ b/fastvideo/dataset/transform.py @@ -40,7 +40,6 @@ def center_crop_th_tw(clip, th, tw, top_crop) -> torch.Tensor: if not _is_tensor_video_clip(clip): raise ValueError("clip should be a 4D torch.tensor") - # import ipdb;ipdb.set_trace() h, w = clip.size(-2), clip.size(-1) tr = th / tw if h / w > tr: @@ -67,7 +66,6 @@ def normalize_video(clip) -> torch.Tensor: _is_tensor_video_clip(clip) if not clip.dtype == torch.uint8: raise TypeError(f"clip tensor should have data type uint8. Got {clip.dtype}") - # return clip.float().permute(3, 0, 1, 2) / 255.0 return clip.float() / 255.0 diff --git a/fastvideo/dataset/utils.py b/fastvideo/dataset/utils.py index 660837fef..5756cc489 100644 --- a/fastvideo/dataset/utils.py +++ b/fastvideo/dataset/utils.py @@ -154,23 +154,18 @@ def collate_latents_embs_masks(batch_to_process, primary-encoder-only: the secondary ByT5 embedding (``text_embedding_2``) and its mask are not read even when the parquet rows contain them. """ - # Initialize tensors to hold padded embeddings and masks all_latents = [] all_embs = [] all_masks = [] caption_text = [] - # Process each row individually for i, row in enumerate(batch_to_process): - # Get tensors from row data = get_torch_tensors_from_row_dict(row, keys, cfg_rate, rng) latents, emb = data["vae_latent"], data["text_embedding"] padded_emb, mask = pad(emb, text_padding_length) - # Store in batch tensors all_latents.append(latents) all_embs.append(padded_emb) all_masks.append(mask) - # TODO(py): remove this once we fix preprocess try: caption_text.append(row["prompt"]) except KeyError: @@ -262,10 +257,8 @@ def collate_rows_from_parquet_schema(rows, # callers, including their configured text encoder length. effective_text_padding_length = text_padding_length - # Initialize containers for different data types batch_data: dict[str, Any] = {} - # Get tensor and metadata field names from schema (fields ending with '_bytes') tensor_fields = [] metadata_fields = [] for field in parquet_schema.names: @@ -297,7 +290,6 @@ def collate_rows_from_parquet_schema(rows, row_drops[row_idx] = ((rng.random() if rng else random.random()) < cfg_rate) return row_drops[row_idx] - # Process each tensor field for tensor_name in tensor_fields: tensor_list = [] shape_key = f"{tensor_name}_shape" @@ -322,7 +314,6 @@ def collate_rows_from_parquet_schema(rows, "text_embedding_bytes and text_embedding_shape.") for row_idx, row in enumerate(rows): - # Get tensor data from row using the existing helper function pattern if shape_key in row and bytes_key in row: shape = row[shape_key] bytes_data = row[bytes_key] @@ -343,19 +334,13 @@ def collate_rows_from_parquet_schema(rows, row.get(dtype_key, "float32"), zero=drop, ) - # if len(data.shape) == 3: - # B, L, D = tensor.shape - # assert B == 1, "Batch size must be 1" - # tensor = tensor.squeeze(0) tensor_list.append(tensor) else: - # Handle missing tensor data tensor_list.append(torch.zeros(0, dtype=torch.bfloat16)) # Stack tensors with special handling for text embeddings if tensor_name in _TEXT_STREAM_MASKS: - # Handle text embeddings with padding padded_tensors = [] attention_masks = [] # Empty entries have to stack against the real ones, so follow @@ -373,7 +358,6 @@ def collate_rows_from_parquet_schema(rows, padded_tensors.append(padded_tensor) attention_masks.append(mask) else: - # Handle empty embeddings - assume default embedding dimension width = (int(tensor.shape[1]) if tensor.dim() > 1 else empty_width) padded_tensors.append(torch.zeros(effective_text_padding_length, width, dtype=empty_dtype)) attention_masks.append(torch.zeros(effective_text_padding_length)) @@ -411,7 +395,6 @@ def collate_rows_from_parquet_schema(rows, f"All tensors in a batch must have compatible shapes. " f"Original error: {e}") from e - # Process metadata fields into info_list info_list = [] for row in rows: info = {} diff --git a/fastvideo/dataset/validation_dataset.py b/fastvideo/dataset/validation_dataset.py index 755df34f6..422c206a3 100644 --- a/fastvideo/dataset/validation_dataset.py +++ b/fastvideo/dataset/validation_dataset.py @@ -20,7 +20,6 @@ class ValidationDataset(IterableDataset): super().__init__() self.filename = pathlib.Path(filename) - # get directory of filename self.dir = os.path.abspath(self.filename.parent) if not self.filename.exists(): @@ -40,17 +39,14 @@ class ValidationDataset(IterableDataset): f"Unsupported file format {self.filename.suffix} for validation dataset. Supported formats are: {_SUPPORTED_FILE_FORMATS}" ) - # Get distributed training info self.global_rank = get_world_rank() self.world_size = get_world_size() self.sp_world_size = get_sp_world_size() self.num_sp_groups = self.world_size // self.sp_world_size - # Convert to list to get total samples self.all_samples = list(data) self.original_total_samples = len(self.all_samples) - # Extend samples to be a multiple of DP degree (num_sp_groups) remainder = self.original_total_samples % self.num_sp_groups if remainder != 0: samples_to_add = self.num_sp_groups - remainder @@ -64,17 +60,13 @@ class ValidationDataset(IterableDataset): self.total_samples = len(self.all_samples) - # Calculate which SP group this rank belongs to self.sp_group_id = self.global_rank // self.sp_world_size - # Now all SP groups will have equal number of samples self.samples_per_sp_group = self.total_samples // self.num_sp_groups - # Calculate start and end indices for this SP group self.start_idx = self.sp_group_id * self.samples_per_sp_group self.end_idx = self.start_idx + self.samples_per_sp_group - # Get samples for this SP group self.sp_group_samples = self.all_samples[self.start_idx:self.end_idx] logger.info( @@ -102,8 +94,6 @@ class ValidationDataset(IterableDataset): # However, since the model specifications use "prompt", we create an alias here. sample["prompt"] = sample["caption"] - # Load image or video if the path is provided - # TODO(aryan): need to handle custom columns here for control conditions sample["image"] = None sample["video"] = None diff --git a/fastvideo/distributed/__init__.py b/fastvideo/distributed/__init__.py index a3fe6c0b8..009d5dc78 100644 --- a/fastvideo/distributed/__init__.py +++ b/fastvideo/distributed/__init__.py @@ -37,6 +37,5 @@ __all__ = [ "get_tp_rank", "get_tp_world_size", - # Get torch device "get_local_torch_device", ] diff --git a/fastvideo/distributed/communication_op.py b/fastvideo/distributed/communication_op.py index eb6d207a7..011ffbb57 100644 --- a/fastvideo/distributed/communication_op.py +++ b/fastvideo/distributed/communication_op.py @@ -24,7 +24,6 @@ def tensor_model_parallel_all_gather(input_: torch.Tensor, dim: int = -1) -> tor return get_tp_group().all_gather(input_, dim) -# TODO: remove model, make it sequence_parallel def sequence_model_parallel_all_to_all_4D(input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1) -> torch.Tensor: @@ -78,7 +77,6 @@ def sequence_model_parallel_shard(input_: torch.Tensor, dim: int = 1) -> tuple[t original_seq_len = input_.shape[dim] - # Compute padding if needed padded_seq_len, padding_amount = compute_padding_for_sp(original_seq_len, sp_world_size) # Pad if necessary @@ -136,7 +134,6 @@ def warmup_sequence_parallel_communication(device: torch.device | None = None) - num_heads = sp_world_size * 4 # Must be divisible by sp_world_size head_dim = 64 - # Create dummy tensor for warmup dummy = torch.zeros(batch_size, seq_len_per_rank, num_heads, head_dim, device=device, dtype=torch.bfloat16) # Warmup pattern 1: scatter heads, gather sequence (before attention) @@ -158,7 +155,6 @@ def warmup_sequence_parallel_communication(device: torch.device | None = None) - # Synchronize to ensure warmup completes device_module.synchronize(device) - # Clean up del dummy, dummy2, dummy3 _sp_warmup_done = True diff --git a/fastvideo/distributed/device_communicators/base_device_communicator.py b/fastvideo/distributed/device_communicators/base_device_communicator.py index 1adfecb1b..1fa65d4a4 100644 --- a/fastvideo/distributed/device_communicators/base_device_communicator.py +++ b/fastvideo/distributed/device_communicators/base_device_communicator.py @@ -246,12 +246,10 @@ class DeviceCommunicatorBase: world_size = self.world_size assert -input_.dim() <= dim < input_.dim(), (f"Invalid dim ({dim}) for input tensor with shape {input_.size()}") if dim < 0: - # Convert negative dim to positive. dim += input_.dim() # Allocate output tensor. gather_list = [torch.empty_like(input_) for _ in range(world_size)] if self.rank_in_group == dst else None - # Gather. torch.distributed.gather(input_, gather_list, dst=self.ranks[dst], group=self.device_group) output_tensor = torch.cat(gather_list, dim=dim) if self.rank_in_group == dst else None return output_tensor diff --git a/fastvideo/distributed/device_communicators/cpu_communicator.py b/fastvideo/distributed/device_communicators/cpu_communicator.py index 363b1ae3c..7bac20c60 100644 --- a/fastvideo/distributed/device_communicators/cpu_communicator.py +++ b/fastvideo/distributed/device_communicators/cpu_communicator.py @@ -42,13 +42,11 @@ class CpuCommunicator(DeviceCommunicatorBase): world_size = self.world_size assert -input_.dim() <= dim < input_.dim(), (f"Invalid dim ({dim}) for input tensor with shape {input_.size()}") if dim < 0: - # Convert negative dim to positive. dim += input_.dim() # Allocate output tensor. gather_list = [torch.empty_like(input_) for _ in range(world_size)] if self.rank_in_group == dst else None - # Gather. self.dist_module.gather(input_, gather_list, dst=self.ranks[dst], group=self.device_group) output_tensor = torch.cat(gather_list, dim=dim) if self.rank_in_group == dst else None @@ -56,7 +54,6 @@ class CpuCommunicator(DeviceCommunicatorBase): def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: if dim < 0: - # Convert negative dim to positive. dim += input_.dim() input_size = input_.size() # NOTE: we have to use concat-style all-gather here, diff --git a/fastvideo/distributed/device_communicators/pyhccl.py b/fastvideo/distributed/device_communicators/pyhccl.py index ed5ecda56..cfb75d872 100644 --- a/fastvideo/distributed/device_communicators/pyhccl.py +++ b/fastvideo/distributed/device_communicators/pyhccl.py @@ -70,16 +70,13 @@ class PyHcclCommunicator: device = torch.device(f"npu:{device}") elif isinstance(device, str): device = torch.device(device) - # now `device` is a `torch.device` object assert isinstance(device, torch.device) self.device = device if self.rank == 0: - # get the unique id from HCCL with torch.npu.device(device): self.unique_id = self.hccl.hcclGetUniqueId() else: - # construct an empty unique id self.unique_id = hcclUniqueId() if not isinstance(group, StatelessProcessGroup): diff --git a/fastvideo/distributed/device_communicators/pynccl.py b/fastvideo/distributed/device_communicators/pynccl.py index 6aef9bb35..b48a3f6e4 100644 --- a/fastvideo/distributed/device_communicators/pynccl.py +++ b/fastvideo/distributed/device_communicators/pynccl.py @@ -68,10 +68,8 @@ class PyNcclCommunicator: logger.info("FastVideo is using nccl==%s", self.nccl.ncclGetVersion()) if self.rank == 0: - # get the unique id from NCCL self.unique_id = self.nccl.ncclGetUniqueId() else: - # construct an empty unique id self.unique_id = ncclUniqueId() if not isinstance(group, StatelessProcessGroup): @@ -88,7 +86,6 @@ class PyNcclCommunicator: device = torch.device(f"cuda:{device}") elif isinstance(device, str): device = torch.device(device) - # now `device` is a `torch.device` object assert isinstance(device, torch.device) self.device = device # nccl communicator and stream will use this device diff --git a/fastvideo/distributed/device_communicators/pynccl_wrapper.py b/fastvideo/distributed/device_communicators/pynccl_wrapper.py index 7f7f74353..21976d306 100644 --- a/fastvideo/distributed/device_communicators/pynccl_wrapper.py +++ b/fastvideo/distributed/device_communicators/pynccl_wrapper.py @@ -22,7 +22,6 @@ # changing the environment variable `FASTVIDEO_NCCL_SO_PATH`, or the `so_file` # variable in the code. -#TODO(will): support FASTVIDEO_NCCL_SO_PATH import ctypes import platform diff --git a/fastvideo/distributed/parallel_state.py b/fastvideo/distributed/parallel_state.py index 4fbb0a1d1..00207ac6a 100644 --- a/fastvideo/distributed/parallel_state.py +++ b/fastvideo/distributed/parallel_state.py @@ -187,7 +187,6 @@ class GroupCoordinator: from fastvideo.platforms import current_platform - # TODO: fix it for other platforms self.device = get_local_torch_device() self.use_device_communicator = use_device_communicator @@ -223,8 +222,6 @@ class GroupCoordinator: from fastvideo.platforms import current_platform - # TODO(will): check if this is needed - # self.use_custom_op_call = current_platform.is_cuda_alike() self.use_custom_op_call = False @property @@ -284,7 +281,6 @@ class GroupCoordinator: else: # For non-CUDA platforms (MPS, CPU), just yield the context without stream management if graph_capture_context is None: - # Create a dummy context for non-CUDA platforms graph_capture_context = GraphCaptureContext(None) yield graph_capture_context @@ -475,16 +471,13 @@ class GroupCoordinator: async_handles = [] for tensor in tensor_list: if tensor.numel() == 0: - # Skip broadcasting empty tensors. continue if tensor.is_cpu: - # use metadata_group for CPU tensors handle = torch.distributed.broadcast(tensor, src=self.ranks[src], group=metadata_group, async_op=True) else: - # use group for GPU tensors handle = torch.distributed.broadcast(tensor, src=self.ranks[src], group=group, async_op=True) async_handles.append(handle) for async_handle in async_handles: @@ -498,17 +491,14 @@ class GroupCoordinator: if isinstance(value, TensorMetadata): tensor = torch.empty(value.size, dtype=value.dtype, device=value.device) if tensor.numel() == 0: - # Skip broadcasting empty tensors. tensor_dict[key] = tensor continue if tensor.is_cpu: - # use metadata_group for CPU tensors handle = torch.distributed.broadcast(tensor, src=self.ranks[src], group=metadata_group, async_op=True) else: - # use group for GPU tensors handle = torch.distributed.broadcast(tensor, src=self.ranks[src], group=group, async_op=True) async_handles.append(handle) tensor_dict[key] = tensor @@ -550,7 +540,6 @@ class GroupCoordinator: self.send_object(metadata_list, dst=dst) for tensor in tensor_list: if tensor.numel() == 0: - # Skip sending empty tensors. continue # send-allgather: send only a slice, then do allgather. @@ -558,10 +547,8 @@ class GroupCoordinator: tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank] if tensor.is_cpu: - # use metadata_group for CPU tensors torch.distributed.send(tensor, dst=self.ranks[dst], group=metadata_group) else: - # use group for GPU tensors torch.distributed.send(tensor, dst=self.ranks[dst], group=group) return None @@ -593,7 +580,6 @@ class GroupCoordinator: if isinstance(value, TensorMetadata): tensor = torch.empty(value.size, dtype=value.dtype, device=value.device) if tensor.numel() == 0: - # Skip broadcasting empty tensors. tensor_dict[key] = tensor continue @@ -605,10 +591,8 @@ class GroupCoordinator: tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank] if tensor.is_cpu: - # use metadata_group for CPU tensors torch.distributed.recv(tensor, src=self.ranks[src], group=metadata_group) else: - # use group for GPU tensors torch.distributed.recv(tensor, src=self.ranks[src], group=group) if use_all_gather: # do the allgather @@ -789,7 +773,6 @@ def init_distributed_environment( else: assert _WORLD.world_size == torch.distributed.get_world_size(), ( "world group already initialized with a different world size") - # Init a group for each node if _NODE is None: init_node_group(local_rank, backend) @@ -832,7 +815,6 @@ def initialize_model_parallel( must still belong to exactly one group. dp_group_ranks: explicit DP group layout, same contract. """ - # Get world size and rank. Ensure some consistencies. assert _WORLD is not None, "world group is not initialized, please call init_distributed_environment first" world_size: int = get_world_size() backend = backend or torch.distributed.get_backend(get_world_group().device_group) @@ -852,7 +834,6 @@ def initialize_model_parallel( use_message_queue_broadcaster=True, group_name="tp") - # Build the sequence model-parallel groups. num_sequence_model_parallel_groups: int = (world_size // sequence_model_parallel_size) global _SP assert _SP is None, ("sequence model parallel group is already initialized") @@ -868,13 +849,11 @@ def initialize_model_parallel( else: # Since SP is incompatible with TP and PP, we can use a simpler group creation logic for i in range(num_sequence_model_parallel_groups): - # Create groups of consecutive ranks ranks = list(range(i * sequence_model_parallel_size, (i + 1) * sequence_model_parallel_size)) group_ranks.append(ranks) _SP = init_model_parallel_group(group_ranks, get_world_group().local_rank, backend, group_name="sp") - # Build the data parallel groups. num_data_parallel_groups: int = sequence_model_parallel_size global _DP assert _DP is None, ("data parallel group is already initialized") @@ -942,7 +921,6 @@ def maybe_init_distributed_environment_and_model_parallel(tp_size: int, sp_group_ranks: list[list[int]] | None = None, dp_group_ranks: list[list[int]] | None = None): if _WORLD is not None and model_parallel_is_initialized(): - # make sure the tp and sp sizes are correct assert get_tp_world_size( ) == tp_size, f"You are trying to initialize model parallel groups with size {tp_size}, but they are already initialized with size {get_tp_world_size()}" assert get_sp_world_size( @@ -968,7 +946,6 @@ def maybe_init_distributed_environment_and_model_parallel(tp_size: int, sp_group_ranks=sp_group_ranks, dp_group_ranks=dp_group_ranks) - # set device if we're on a CUDA/NPU platform from fastvideo.platforms import current_platform if current_platform.is_cuda_alike() or current_platform.is_npu(): device_type = current_platform.device_type @@ -1091,7 +1068,6 @@ def get_same_node_ranks(pg: ProcessGroup | StatelessProcessGroup, source_rank: i rank = torch.distributed.get_rank(group=pg) world_size = torch.distributed.get_world_size(group=pg) - # global ranks of the processes in the group ranks = torch.distributed.get_process_group_ranks(pg) else: rank = pg.rank @@ -1107,7 +1083,6 @@ def get_same_node_ranks(pg: ProcessGroup | StatelessProcessGroup, source_rank: i try: with contextlib.suppress(OSError): if rank == source_rank: - # create a shared memory segment shm = shared_memory.SharedMemory(create=True, size=128) shm.buf[:len(magic_message)] = magic_message if isinstance(pg, ProcessGroup): @@ -1141,7 +1116,6 @@ def get_same_node_ranks(pg: ProcessGroup | StatelessProcessGroup, source_rank: i else: pg.barrier() - # clean up the shared memory segment with contextlib.suppress(OSError): if rank == source_rank and shm: shm.unlink() @@ -1201,23 +1175,19 @@ def initialize_tensor_parallel_group(tensor_model_parallel_size: int = 1, output1 = model1(input1) ``` """ - # Get world size and rank. Ensure some consistencies. assert torch.distributed.is_initialized() world_size: int = torch.distributed.get_world_size() backend = backend or torch.distributed.get_backend(get_world_group().device_group) - # Ensure the world size is compatible with the parallelism configuration assert world_size % tensor_model_parallel_size == 0, \ f"World size ({world_size}) must be divisible by tensor_model_parallel_size ({tensor_model_parallel_size})" - # Build the tensor model-parallel groups. num_tensor_model_parallel_groups: int = (world_size // tensor_model_parallel_size) tp_group_ranks = [] for i in range(num_tensor_model_parallel_groups): ranks = list(range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size)) tp_group_ranks.append(ranks) - # Create TP group coordinator with a unique name group_name = f"tp_{group_name_suffix}" if group_name_suffix else "tp" tp_group = init_model_parallel_group(tp_group_ranks, get_world_group().local_rank, @@ -1260,25 +1230,20 @@ def initialize_sequence_parallel_group(sequence_model_parallel_size: int = 1, output2 = model2(input2) ``` """ - # Get world size and rank. Ensure some consistencies. assert torch.distributed.is_initialized() world_size: int = torch.distributed.get_world_size() backend = backend or torch.distributed.get_backend(get_world_group().device_group) - # Ensure the world size is compatible with the parallelism configuration assert world_size % sequence_model_parallel_size == 0, \ f"World size ({world_size}) must be divisible by sequence_model_parallel_size ({sequence_model_parallel_size})" - # Build the sequence model-parallel groups. num_sequence_model_parallel_groups: int = (world_size // sequence_model_parallel_size) sp_group_ranks = [] for i in range(num_sequence_model_parallel_groups): - # Create groups of consecutive ranks ranks = list(range(i * sequence_model_parallel_size, (i + 1) * sequence_model_parallel_size)) sp_group_ranks.append(ranks) - # Create SP group coordinator with a unique name group_name = f"sp_{group_name_suffix}" if group_name_suffix else "sp" sp_group = init_model_parallel_group(sp_group_ranks, get_world_group().local_rank, backend, group_name=group_name) diff --git a/fastvideo/distributed/utils.py b/fastvideo/distributed/utils.py index 5e70e9cc4..f439205c8 100644 --- a/fastvideo/distributed/utils.py +++ b/fastvideo/distributed/utils.py @@ -48,7 +48,6 @@ def split_tensor_along_last_dim( Returns: A list of Tensors """ - # Get the size and dimension. last_dim = tensor.dim() - 1 last_dim_size = divide(tensor.size()[last_dim], num_partitions) # Split. @@ -105,11 +104,9 @@ def pad_sequence_tensor( padding_amount = target_seq_len - current_seq_len - # Create padding shape pad_shape = list(tensor.shape) pad_shape[seq_dim] = padding_amount - # Create padding tensor padding = torch.full( pad_shape, pad_value, @@ -139,7 +136,6 @@ def unpad_sequence_tensor( Returns: Tensor: Unpadded tensor """ - # Use slice to remove padding indices = [slice(None)] * tensor.dim() indices[seq_dim] = slice(0, original_seq_len) @@ -184,7 +180,6 @@ class StatelessProcessGroup: def expire_data(self) -> None: """Expire data that is older than `data_expiration_seconds` seconds.""" while self.entries: - # check the oldest entry key, timestamp = self.entries[0] if time.perf_counter() - timestamp > self.data_expiration_seconds: self.store.delete_key(key) diff --git a/fastvideo/entrypoints/cli/bench_serving.py b/fastvideo/entrypoints/cli/bench_serving.py index 33054eb51..756ace686 100644 --- a/fastvideo/entrypoints/cli/bench_serving.py +++ b/fastvideo/entrypoints/cli/bench_serving.py @@ -338,7 +338,6 @@ async def async_request_image_sglang( # Check if we need to use multipart (for image edits with input images) if input.image_paths and len(input.image_paths) > 0: - # Use multipart/form-data for image edits data = aiohttp.FormData() data.add_field("model", input.model) data.add_field("prompt", input.prompt) @@ -351,7 +350,6 @@ async def async_request_image_sglang( for key, value in input.extra_body.items(): data.add_field(key, str(value)) - # Add image file(s) for idx, img_path in enumerate(input.image_paths): if os.path.exists(img_path): async with aiofiles.open(img_path, "rb") as f: @@ -384,7 +382,6 @@ async def async_request_image_sglang( output.error = str(e) output.success = False else: - # Use JSON for text-to-image generation payload = { "model": input.model, "prompt": input.prompt, @@ -432,7 +429,6 @@ async def async_request_video_sglang( job_id = None # Check if we need to upload images (Multipart) or just send JSON if input.image_paths and len(input.image_paths) > 0: - # Use multipart/form-data data = aiohttp.FormData() data.add_field("model", input.model) data.add_field("prompt", input.prompt) @@ -440,7 +436,6 @@ async def async_request_video_sglang( if input.width and input.height: data.add_field("size", f"{input.width}x{input.height}") - # Add extra body fields to form data if possible if input.extra_body: data.add_field("extra_body", json.dumps(input.extra_body)) @@ -449,7 +444,6 @@ async def async_request_video_sglang( if input.fps: data.add_field("fps", str(input.fps)) - # Add image file img_path = input.image_paths[0] if os.path.exists(img_path): async with aiofiles.open(img_path, "rb") as f: @@ -486,7 +480,6 @@ async def async_request_video_sglang( return output else: - # Use JSON payload: dict[str, Any] = { "model": input.model, "prompt": input.prompt, @@ -635,14 +628,12 @@ async def wait_for_service(base_url: str, timeout: int = 1200) -> None: async def benchmark(args: argparse.Namespace) -> None: from huggingface_hub import model_info - # Construct base_url if not provided if args.base_url is None: args.base_url = f"http://{args.host}:{args.port}" # Wait for service await wait_for_service(args.base_url) - # Fetch model info try: async with aiohttp.ClientSession() as session, session.get( f"{args.base_url}/v1/model_info", @@ -695,7 +686,6 @@ async def benchmark(args: argparse.Namespace) -> None: # Limit concurrency semaphore = asyncio.Semaphore(args.max_concurrency) if args.max_concurrency is not None else None - # Run benchmark pbar = tqdm(total=len(requests_list), disable=args.disable_tqdm) async with aiohttp.ClientSession() as session: @@ -717,17 +707,14 @@ async def benchmark(args: argparse.Namespace) -> None: pbar.close() - # Calculate metrics metrics = calculate_metrics(outputs, total_duration) print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=60, c="=")) - # Section 1: Configuration print("{:<40} {:<15}".format("Task:", task_name)) print("{:<40} {:<15}".format("Model:", args.model)) print("{:<40} {:<15}".format("Dataset:", args.dataset)) - # Section 2: Execution & Traffic print(f"{'-' * 50}") print("{:<40} {:<15.2f}".format("Benchmark duration (s):", metrics["duration"])) print("{:<40} {:<15}".format("Request rate:", str(args.request_rate))) @@ -737,7 +724,6 @@ async def benchmark(args: argparse.Namespace) -> None: )) print("{:<40} {}/{:<15}".format("Successful requests:", metrics["completed_requests"], len(requests_list))) - # Section 3: Performance Metrics print(f"{'-' * 50}") print("{:<40} {:<15.2f}".format("Request throughput (req/s):", metrics["throughput_qps"])) diff --git a/fastvideo/entrypoints/openai/api_server.py b/fastvideo/entrypoints/openai/api_server.py index dfe20c476..0665b6ec6 100644 --- a/fastvideo/entrypoints/openai/api_server.py +++ b/fastvideo/entrypoints/openai/api_server.py @@ -170,7 +170,6 @@ def create_app( }}, ) - # Import and mount routers from fastvideo.entrypoints.openai.common_api import router as common_router from fastvideo.entrypoints.openai.image_api import router as image_router from fastvideo.entrypoints.openai.playground import router as playground_router @@ -209,7 +208,6 @@ def _parse_args() -> tuple[FastVideoArgs, str, int, str]: port = args.port output_dir = args.output_dir - # Build FastVideoArgs from the remaining CLI args excluded = {"host", "port", "output_dir", "subparser", "config", "dispatch_function"} cli_kwargs = {k: v for k, v in vars(args).items() if k not in excluded and v is not None} fastvideo_args = FastVideoArgs.from_kwargs(**cli_kwargs) diff --git a/fastvideo/entrypoints/openai/image_api.py b/fastvideo/entrypoints/openai/image_api.py index d04910490..93b2dea60 100644 --- a/fastvideo/entrypoints/openai/image_api.py +++ b/fastvideo/entrypoints/openai/image_api.py @@ -265,7 +265,6 @@ async def edits( enable_teacache=enable_teacache, ) - # Save input images uploads_dir = os.path.join(get_output_dir(), "uploads") os.makedirs(uploads_dir, exist_ok=True) image_list = merge_image_input_list(images, urls) diff --git a/fastvideo/entrypoints/streaming/server.py b/fastvideo/entrypoints/streaming/server.py index fcfb94bf7..d03743472 100644 --- a/fastvideo/entrypoints/streaming/server.py +++ b/fastvideo/entrypoints/streaming/server.py @@ -385,8 +385,6 @@ async def _run_segment( )) start = time.perf_counter() - # TODO: pool.run() runs to completion even if the client disconnects - # mid-segment. Real cancellation needs the generate_async API. try: result = await state.pool.run(session.id, request) except Exception as exc: diff --git a/fastvideo/entrypoints/streaming_generator.py b/fastvideo/entrypoints/streaming_generator.py index 8a7913057..6e23fd971 100644 --- a/fastvideo/entrypoints/streaming_generator.py +++ b/fastvideo/entrypoints/streaming_generator.py @@ -46,7 +46,6 @@ class IncrementalVideoWriter: frames_copy = [f.copy() for f in frames] self._pending_main = self._executor.submit(self._write_frames, frames_copy) - # Write block file if block_dir is set block_future = None if self._block_dir: self._block_idx += 1 @@ -132,7 +131,6 @@ class StreamingVideoGenerator(VideoGenerator): self.writer = None self.executor.execute_streaming_clear() - # Handle batch processing from text file if self.sampling_param is None: self.sampling_param = SamplingParam.from_pretrained(self.fastvideo_args.model_path) @@ -144,7 +142,6 @@ class StreamingVideoGenerator(VideoGenerator): if "output_path" in kwargs: output_path = self._prepare_output_path(kwargs["output_path"], prompt) - # Create block directory for individual block files block_dir = output_path.replace(".mp4", "") os.makedirs(block_dir, exist_ok=True) self.block_dir = block_dir diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index dc52f444c..8be4cd670 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -316,8 +316,6 @@ class VideoGenerator: Returns: The created video generator """ - # Initialize distributed environment if needed - # initialize_distributed_and_parallelism(fastvideo_args) executor_class = Executor.get_class(fastvideo_args, allow_external_launcher=True) return cls( @@ -614,14 +612,12 @@ class VideoGenerator: validate_request_batch_extra(kwargs, model_path=fastvideo_args.model_path) - # Handle batch processing from text file if sampling_param is None: sampling_param = SamplingParam.from_pretrained(fastvideo_args.model_path) if self.executor.uses_spmd_execution: prompt = self.executor.broadcast_from_output_rank(prompt if self.executor.is_output_rank else None) - # Add action control inputs to kwargs if provided if mouse_cond is not None: kwargs['mouse_cond'] = mouse_cond if keyboard_cond is not None: @@ -666,7 +662,6 @@ class VideoGenerator: for i, batch_prompt in enumerate(prompts): logger.info("Processing prompt %d/%d: %s...", i + 1, len(prompts), batch_prompt[:100]) try: - # Generate video for this prompt using the same logic below output_path = self._run_on_output_rank( partial(self._prepare_output_path, sampling_param.output_path, @@ -683,7 +678,6 @@ class VideoGenerator: **kwargs, ) - # Add prompt info to result result["prompt_index"] = i result["prompt"] = batch_prompt @@ -774,7 +768,6 @@ class VideoGenerator: target_ext = ".mp4" def _sanitize_filename_component(name: str) -> str: - # Remove characters invalid on common filesystems, strip spaces/dots sanitized = re.sub(r'[\\/:*?"<>|]', '', name) sanitized = sanitized.strip().strip('.') sanitized = re.sub(r'\s+', ' ', sanitized) @@ -851,7 +844,6 @@ class VideoGenerator: if fastvideo_args is None: fastvideo_args = self.fastvideo_args - # Validate inputs if not isinstance(prompt, str): raise TypeError(f"`prompt` must be a string, but got {type(prompt)}") prompt = prompt.strip() @@ -859,21 +851,17 @@ class VideoGenerator: output_path = kwargs["output_path"] prompt_embeds = kwargs.get("prompt_embeds") sampling_param.prompt = prompt - # Process negative prompt if sampling_param.negative_prompt is not None: sampling_param.negative_prompt = sampling_param.negative_prompt.strip() - # Validate dimensions if (sampling_param.height <= 0 or sampling_param.width <= 0 or sampling_param.num_frames <= 0): raise ValueError(f"Height, width, and num_frames must be positive integers, got " f"height={sampling_param.height}, width={sampling_param.width}, " f"num_frames={sampling_param.num_frames}") - # Calculate sizes target_height = align_to(sampling_param.height, 16) target_width = align_to(sampling_param.width, 16) - # Calculate latent sizes latents_size = [(sampling_param.num_frames - 1) // 4 + 1, sampling_param.height // 8, sampling_param.width // 8] n_tokens = latents_size[0] * latents_size[1] * latents_size[2] @@ -897,7 +885,6 @@ class VideoGenerator: """ # type: ignore[attr-defined] logger.info(debug_str) - # Prepare batch batch = ForwardBatch( **shallow_asdict(sampling_param), eta=0.0, @@ -924,7 +911,6 @@ class VideoGenerator: batch.return_trajectory_decoded = False batch.return_continuation_state = False - # Run inference start_time = time.perf_counter() # Execute forward pass in a new thread for non-blocking tensor @@ -1476,14 +1462,11 @@ class VideoGenerator: num_channels = cls._write_pcm_wav(wav_path, audio, sample_rate) layout = "stereo" if num_channels == 2 else "mono" - # Open input video and audio input_video = av.open(video_path) input_audio = av.open(wav_path) - # Create output with both streams output = av.open(out_path, mode="w") - # Add video stream (copy codec from input) in_video_stream = input_video.streams.video[0] out_video_stream = output.add_stream( codec_name=in_video_stream.codec_context.name, @@ -1493,7 +1476,6 @@ class VideoGenerator: out_video_stream.height = in_video_stream.height out_video_stream.pix_fmt = in_video_stream.pix_fmt - # Add audio stream (AAC) out_audio_stream = output.add_stream("aac", rate=sample_rate) out_audio_stream.layout = layout diff --git a/fastvideo/eval/io/inputs.py b/fastvideo/eval/io/inputs.py index 95e2aeccf..fbf9dfc9b 100644 --- a/fastvideo/eval/io/inputs.py +++ b/fastvideo/eval/io/inputs.py @@ -262,9 +262,6 @@ def samples_from_manifest( return samples -# ------------------------------------------------------------------ -# Internal helpers -# ------------------------------------------------------------------ def _expand(spec: PathSpec, exts: tuple[str, ...]) -> list[Path]: @@ -340,7 +337,6 @@ def _attach_extracted_audio( else: out_dir = Path(tempfile.mkdtemp(prefix="fv_extracted_audio_")) - # Build the work list — (sample_idx, dest_audio_key, source_path) work: list[tuple[int, str, str]] = [] for i, s in enumerate(samples): for vkey, akey in ( diff --git a/fastvideo/eval/metrics/optical_flow/synthetic_optical_flow/_thirdperson.py b/fastvideo/eval/metrics/optical_flow/synthetic_optical_flow/_thirdperson.py index 4f372784a..d296e9b09 100644 --- a/fastvideo/eval/metrics/optical_flow/synthetic_optical_flow/_thirdperson.py +++ b/fastvideo/eval/metrics/optical_flow/synthetic_optical_flow/_thirdperson.py @@ -192,9 +192,7 @@ class ThirdPersonFlowGenerator: return [self.generate_flow(kb[i], mo[i]) for i in range(n)] -# --------------------------------------------------------------------------- # Linear-features form (used by the calibration fitter). -# --------------------------------------------------------------------------- # Number of free params we fit. Order is fixed and shared with the fitter. PARAM_NAMES = ( diff --git a/fastvideo/eval/metrics/vbench/aesthetic_quality/metric.py b/fastvideo/eval/metrics/vbench/aesthetic_quality/metric.py index d8be74702..5700e6548 100644 --- a/fastvideo/eval/metrics/vbench/aesthetic_quality/metric.py +++ b/fastvideo/eval/metrics/vbench/aesthetic_quality/metric.py @@ -68,7 +68,6 @@ class AestheticQualityMetric(BaseMetric): ) self._clip_model.eval() - # Load LAION aesthetic head ckpt_path = ensure_checkpoint( "sa_0_4_vit_l_14_linear.pth", source=_AESTHETIC_URL, diff --git a/fastvideo/eval/metrics/vbench/color/metric.py b/fastvideo/eval/metrics/vbench/color/metric.py index 250f5ddce..77759002b 100644 --- a/fastvideo/eval/metrics/vbench/color/metric.py +++ b/fastvideo/eval/metrics/vbench/color/metric.py @@ -61,7 +61,6 @@ class ColorMetric(BaseMetric): prompt = sample.get("text_prompt") or "" color_key = aux["color"] - # Parse object name: remove "a ", "an ", and the color word object_key = prompt.replace("a ", "").replace("an ", "").replace(color_key, "").strip() frames_np = prepare_frames(video) diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 607660f36..1e963492e 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -969,10 +969,8 @@ class FastVideoArgs: type=str, help="Path to safetensors file for initial weight loading") - # Add pipeline configuration arguments PipelineConfig.add_cli_args(parser) - # Add preprocessing configuration arguments PreprocessConfig.add_cli_args(parser) return parser @@ -980,10 +978,8 @@ class FastVideoArgs: @classmethod def from_cli_args(cls, args: argparse.Namespace) -> "FastVideoArgs": provided_args = clean_cli_args(args) - # Get all fields from the dataclass attrs = [attr.name for attr in dataclasses.fields(cls)] - # Create a dictionary of attribute values, with defaults for missing attributes kwargs: dict[str, Any] = {} for attr in attrs: if attr == 'pipeline_config': @@ -993,11 +989,9 @@ class FastVideoArgs: preprocess_config = PreprocessConfig.from_kwargs(provided_args) kwargs['preprocess_config'] = preprocess_config elif attr == 'mode': - # Convert string to ExecutionMode enum mode_value = getattr(args, attr, FastVideoArgs.mode.value) kwargs['mode'] = ExecutionMode.from_string(mode_value) if isinstance(mode_value, str) else mode_value elif attr == 'torch_compile_kwargs': - # Parse JSON string for torch.compile kwargs torch_compile_kwargs_str = getattr(args, 'torch_compile_kwargs', None) if torch_compile_kwargs_str: try: @@ -1008,17 +1002,13 @@ class FastVideoArgs: else: kwargs['torch_compile_kwargs'] = {} elif attr == 'workload_type': - # Convert string to WorkloadType enum workload_type_value = getattr(args, 'workload_type', FastVideoArgs.workload_type.value) kwargs['workload_type'] = WorkloadType.from_string(workload_type_value) if isinstance( workload_type_value, str) else workload_type_value - # Use getattr with default value from the dataclass for potentially missing attributes else: - # Get the field to check if it has a default_factory field = dataclasses.fields(cls)[next(i for i, f in enumerate(dataclasses.fields(cls)) if f.name == attr)] if field.default_factory is not dataclasses.MISSING: - # Use the default_factory to create the default value default_value = field.default_factory() else: default_value = getattr(cls, attr, None) @@ -1029,11 +1019,9 @@ class FastVideoArgs: @classmethod def from_kwargs(cls, **kwargs: Any) -> "FastVideoArgs": - # Convert mode string to enum if necessary if 'mode' in kwargs and isinstance(kwargs['mode'], str): kwargs['mode'] = ExecutionMode.from_string(kwargs['mode']) - # Convert workload_type string to enum if necessary if 'workload_type' in kwargs and isinstance(kwargs['workload_type'], str): kwargs['workload_type'] = WorkloadType.from_string(kwargs['workload_type']) @@ -1052,11 +1040,9 @@ class FastVideoArgs: if self.dist_timeout is not None and self.dist_timeout <= 0: raise ValueError(f"dist_timeout must be greater than zero seconds, got {self.dist_timeout}") - # Validate mode and inference_mode consistency assert isinstance(self.mode, ExecutionMode), f"Mode must be an ExecutionMode enum, got {type(self.mode)}" assert self.mode in ExecutionMode.choices(), f"Invalid execution mode: {self.mode}" - # Validate workload type assert isinstance(self.workload_type, WorkloadType), f"Workload type must be a WorkloadType enum, got {type(self.workload_type)}" assert self.workload_type in WorkloadType.choices(), f"Invalid workload type: {self.workload_type}" @@ -1120,7 +1106,6 @@ class FastVideoArgs: self.pipeline_config.check_pipeline_config() - # Add preprocessing config validation if needed if self.mode == ExecutionMode.PREPROCESS: if self.preprocess_config is None: raise ValueError("preprocess_config is not set in FastVideoArgs when mode is PREPROCESS") @@ -1255,7 +1240,6 @@ def get_current_fastvideo_args() -> FastVideoArgs: # in ci, usually when we test custom ops/modules directly, # we don't set the fastvideo config. In that case, we set a default # config. - # TODO(will): may need to handle this for CI. raise ValueError("Current fastvideo args is not set.") return _current_fastvideo_args @@ -1434,27 +1418,21 @@ class TrainingArgs(FastVideoArgs): @classmethod def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs": provided_args = clean_cli_args(args) - # Get all fields from the dataclass attrs = [attr.name for attr in dataclasses.fields(cls)] logger.info(provided_args) - # Create a dictionary of attribute values, with defaults for missing attributes kwargs: dict[str, Any] = {} for attr in attrs: if attr == 'pipeline_config': pipeline_config = PipelineConfig.from_kwargs(provided_args) kwargs[attr] = pipeline_config elif attr == 'mode': - # Convert string to ExecutionMode enum mode_value = getattr(args, attr, ExecutionMode.FINETUNING.value) kwargs[attr] = ExecutionMode.from_string(mode_value) if isinstance(mode_value, str) else mode_value elif attr == 'workload_type': - # Convert string to WorkloadType enum workload_type_value = getattr(args, 'workload_type', WorkloadType.T2V.value) kwargs[attr] = WorkloadType.from_string(workload_type_value) if isinstance(workload_type_value, str) else workload_type_value - # Use getattr with default value from the dataclass for potentially missing attributes else: - # Get the field to check its default value field = dataclasses.fields(cls)[next(i for i, f in enumerate(dataclasses.fields(cls)) if f.name == attr)] @@ -1462,7 +1440,6 @@ class TrainingArgs(FastVideoArgs): if hasattr(args, attr): value = getattr(args, attr) else: - # Use the field's default value if field.default_factory is not dataclasses.MISSING: value = field.default_factory() elif field.default is not dataclasses.MISSING: diff --git a/fastvideo/forward_context.py b/fastvideo/forward_context.py index c5fac4619..69f1157e2 100644 --- a/fastvideo/forward_context.py +++ b/fastvideo/forward_context.py @@ -17,8 +17,6 @@ if TYPE_CHECKING: logger = init_logger(__name__) -# TODO(will): check if this is needed -# track_batchsize: bool = envs.FASTVIDEO_LOG_BATCHSIZE_INTERVAL >= 0 track_batchsize: bool = False last_logging_time: float = 0 forward_start_time: float = 0 @@ -27,14 +25,9 @@ batchsize_logging_interval: float = 1000 batchsize_forward_time: defaultdict = defaultdict(list) -# @dataclass class ForwardContext: current_timestep: int - # TODO(will): check this arg - # copy from vllm_config.compilation_config.static_forward_context - # attn_layers: Dict[str, Any] - # TODO: extend to support per-layer dynamic forward context attn_metadata: "AttentionMetadata" # set dynamically for each forward pass forward_batch: Optional["ForwardBatch"] = None @@ -49,7 +42,6 @@ def get_forward_context() -> "ForwardContext": return _forward_context -# TODO(will): finalize the interface @contextmanager def set_forward_context(current_timestep, attn_metadata, forward_batch: Optional["ForwardBatch"] = None): """A context manager that stores the current forward context, diff --git a/fastvideo/hooks/layerwise_offload.py b/fastvideo/hooks/layerwise_offload.py index 211ea0f3e..9460dd9a2 100644 --- a/fastvideo/hooks/layerwise_offload.py +++ b/fastvideo/hooks/layerwise_offload.py @@ -81,7 +81,6 @@ class LayerwiseOffloadState: @torch.compiler.disable def wait_and_replace_params(self): torch.cuda.current_stream().wait_stream(self.async_copy_stream) - # now gpu_named_parameters are ready for name, param in _offload_tensors(self.module_ref, self.cpu_named_parameters): if not self._will_offload(name): continue @@ -163,7 +162,6 @@ class LayerwiseOffloadHook(ForwardHook): @contextmanager def mutate_params_scope(self): try: - # load params to GPU and keep them there self.state.wait_and_replace_params() # pyright: ignore yield finally: diff --git a/fastvideo/image_processor.py b/fastvideo/image_processor.py index 76b4c90bb..5acd2285d 100644 --- a/fastvideo/image_processor.py +++ b/fastvideo/image_processor.py @@ -43,7 +43,6 @@ class ImageProcessor: torch.Tensor: Normalized tensor of shape (1, 3, height, width) or (1, 1, height, width) for grayscale, with values in range [-1, 1]. """ - # Handle different input types if isinstance(image, PIL.Image.Image): return self._preprocess_pil(image, height, width) elif isinstance(image, np.ndarray): @@ -111,7 +110,6 @@ class ImageProcessor: pil_image = pil_image.resize((width, height), resample=PIL.Image.Resampling.LANCZOS) image_np = np.array(pil_image, dtype=np.float32) / 255.0 - # Ensure 3D shape if image_np.ndim == 2: image_np = np.expand_dims(image_np, axis=-1) @@ -168,7 +166,6 @@ class ImageProcessor: Returns: torch.Tensor: Shape (1, C, H, W) or (1, 1, H, W) with values in [-1, 1] """ - # Convert to tensor if image_np.ndim == 2: # (H, W) - grayscale tensor = torch.from_numpy(image_np).unsqueeze(0).unsqueeze(0) # (1, 1, H, W) elif image_np.ndim == 3: # (H, W, C) @@ -176,7 +173,6 @@ class ImageProcessor: else: raise ValueError(f"Expected 2D or 3D array, got {image_np.ndim}D") - # Normalize to [-1, 1] tensor = 2.0 * tensor - 1.0 return tensor diff --git a/fastvideo/layers/activation.py b/fastvideo/layers/activation.py index 7cd535e4b..0fa40a6b2 100644 --- a/fastvideo/layers/activation.py +++ b/fastvideo/layers/activation.py @@ -7,7 +7,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -# TODO (will): remove this dependency from fastvideo.layers.custom_op import CustomOp diff --git a/fastvideo/layers/fp4linear.py b/fastvideo/layers/fp4linear.py index 04ab6bf3f..0a1841a75 100644 --- a/fastvideo/layers/fp4linear.py +++ b/fastvideo/layers/fp4linear.py @@ -36,7 +36,6 @@ class _LinearFWD4BWD16Fn(torch.autograd.Function): def forward(ctx, x, weight, bias, backend="cutlass", block_size=16, use_128x4_sf_layout=True): flashinfer_mod = _require_flashinfer() - # assert activation dtype if x.dtype not in (torch.float16, torch.bfloat16): x = x.to(dtype=torch.bfloat16) diff --git a/fastvideo/layers/fp8linear.py b/fastvideo/layers/fp8linear.py index 05b8e464f..8c9011a18 100644 --- a/fastvideo/layers/fp8linear.py +++ b/fastvideo/layers/fp8linear.py @@ -48,7 +48,6 @@ class _LinearFWD8BWD16Fn(torch.autograd.Function): @staticmethod def forward(ctx, x, weight, bias, granularity="tensor"): - # assert/normalize activation dtype if x.dtype not in (torch.float16, torch.bfloat16): x = x.to(dtype=torch.bfloat16) @@ -96,7 +95,6 @@ class _LinearFWD8BWD16Fn(torch.autograd.Function): if bias_cast is not None: out2d = out2d + bias_cast - # save tensors for backward (keep original dtypes) ctx.save_for_backward(x2d, weight, bias) ctx.n = n ctx.orig_shape = orig_shape diff --git a/fastvideo/layers/layernorm.py b/fastvideo/layers/layernorm.py index 553b50779..1a13140f9 100644 --- a/fastvideo/layers/layernorm.py +++ b/fastvideo/layers/layernorm.py @@ -213,7 +213,6 @@ class ScaleResidualLayerNormScaleShift(nn.Module): raise ValueError(f"Gate type {type(gate)} not supported") # residual_output.shape: [batch_size, seq_len, inner_dim] - # Apply normalization normalized = self.norm(residual_output) # Identity modulation: skip the scale/shift passes entirely. @@ -224,7 +223,6 @@ class ScaleResidualLayerNormScaleShift(nn.Module): scale = scale.to(normalized.dtype) shift = shift.to(normalized.dtype) - # Apply scale and shift if isinstance(scale, torch.Tensor) and scale.dim() == 4: # scale.shape: [batch_size, num_frames, 1, inner_dim] # shift.shape: [batch_size, num_frames, 1, inner_dim] diff --git a/fastvideo/layers/linear.py b/fastvideo/layers/linear.py index 45db7d182..dadbb60ec 100644 --- a/fastvideo/layers/linear.py +++ b/fastvideo/layers/linear.py @@ -635,7 +635,6 @@ class MergedColumnParallelLinear(ColumnParallelLinear): elif type(param) in (RowvLLMParameter, BasevLLMParameter): param.load_merged_column_weight(loaded_weight=loaded_weight) return - # TODO: @dsikka - move to parameter.py self._load_fused_module_from_checkpoint(param, loaded_weight) return @@ -645,20 +644,6 @@ class MergedColumnParallelLinear(ColumnParallelLinear): if isinstance(param, BlockQuantScaleParameter): raise NotImplementedError("FP8 is not implemented yet") - # FIXME(will): add fp8 support - # from vllm.model_executor.layers.quantization.fp8 import ( - # Fp8LinearMethod, Fp8MoEMethod) - # assert self.quant_method is not None - # assert isinstance(self.quant_method, - # (Fp8LinearMethod, Fp8MoEMethod)) - # weight_block_size = self.quant_method.quant_config.weight_block_size - # assert weight_block_size is not None - # block_n, _ = weight_block_size[0], weight_block_size[1] - # shard_offset = ( - # (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) // - # block_n) // tp_size - # shard_size = ((self.output_sizes[loaded_shard_id] + block_n - 1) // - # block_n // tp_size) else: shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size shard_size = self.output_sizes[loaded_shard_id] // tp_size @@ -810,7 +795,6 @@ class QKVParallelLinear(ColumnParallelLinear): elif type(param) in (RowvLLMParameter, BasevLLMParameter): param.load_qkv_weight(loaded_weight=loaded_weight) return - # TODO: @dsikka - move to parameter.py self._load_fused_module_from_checkpoint(param, loaded_weight) return diff --git a/fastvideo/layers/lora/linear.py b/fastvideo/layers/lora/linear.py index a65ef28c0..3c4ef4c1b 100644 --- a/fastvideo/layers/lora/linear.py +++ b/fastvideo/layers/lora/linear.py @@ -217,7 +217,6 @@ class BaseLayerWithLoRA(nn.Module): self.lora_B = torch.nn.Parameter(B) self.disable_lora = False - # Store rank and alpha directly rank = A.shape[0] # rank is the first dimension of A self.lora_rank = rank self.lora_alpha = lora_alpha if lora_alpha is not None else rank @@ -288,7 +287,6 @@ class BaseLayerWithLoRA(nn.Module): return (weight.to(accumulate_dtype) + (lora_B @ lora_A) * (scale * self.lora_strength)).to(weight.dtype) @torch.no_grad() - # @torch.compile(dynamic=True) def unmerge_lora_weights(self, *, if_merged_only: bool = True) -> None: if self.disable_lora or self.side_branch: return diff --git a/fastvideo/layers/quantization/__init__.py b/fastvideo/layers/quantization/__init__.py index 83f5076f1..34f5fb88a 100644 --- a/fastvideo/layers/quantization/__init__.py +++ b/fastvideo/layers/quantization/__init__.py @@ -76,7 +76,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: "nvfp4_qat_train": NVFP4QATTrainConfig, "fp8_qat_train": FP8QATTrainConfig, } - # Update the `method_to_config` with customized quantization methods. method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG) return method_to_config[quantization] diff --git a/fastvideo/layers/quantization/absmax_fp8.py b/fastvideo/layers/quantization/absmax_fp8.py index 3bbf62d77..58aefe58a 100644 --- a/fastvideo/layers/quantization/absmax_fp8.py +++ b/fastvideo/layers/quantization/absmax_fp8.py @@ -83,7 +83,6 @@ class AbsMaxFP8MergedParameter(nn.Parameter): # MergedColumnParallelLinear case tp_size = get_tp_world_size() if tp_size > 1: - # TODO: support this case raise NotImplementedError( "AbsMaxFP8MergedParameter with integer share_id is not supported in tensor parallelism greater than 1 yet." ) diff --git a/fastvideo/layers/rotary_embedding.py b/fastvideo/layers/rotary_embedding.py index cd8cf59e5..6c7ab8c43 100644 --- a/fastvideo/layers/rotary_embedding.py +++ b/fastvideo/layers/rotary_embedding.py @@ -123,7 +123,6 @@ def _apply_rotary_emb( head_size = x.shape[-1] rope_dim = cos.shape[-1] - # Check if cos/sin are full head_dim (rotate_half style) or half (traditional style) if rope_dim == head_size: # Full head_dim - use rotate_half style (HunyuanVideo) # x * cos + rotate_half(x) * sin @@ -258,17 +257,14 @@ def get_meshgrid_nd(start: int | tuple[int, ...], *args: int | tuple[int, ...], grid (np.ndarray): [dim, ...] """ if len(args) == 0: - # start is grid_size num = _to_tuple(start, dim=dim) start = (0, ) * dim stop = num elif len(args) == 1: - # start is start, args[0] is stop, step is 1 start = _to_tuple(start, dim=dim) stop = _to_tuple(args[0], dim=dim) num = tuple(stop[i] - start[i] for i in range(dim)) elif len(args) == 2: - # start is start, args[0] is stop, args[1] is num start = _to_tuple(start, dim=dim) # Left-Top eg: 12,0 stop = _to_tuple(args[0], dim=dim) # Right-Bottom eg: 20,32 num = _to_tuple(args[1], dim=dim) # Target Size eg: 32,124 @@ -381,35 +377,27 @@ def get_nd_rotary_pos_embed( Returns: Tuple[torch.Tensor, torch.Tensor]: (cos, sin) tensors of shape [HW, D] if use_real, [HW, D/2] otherwise """ - # Get the full grid full_grid = get_meshgrid_nd(start, *args, dim=len(rope_dim_list)) # [3, W, H, D] / [2, W, H] if start_frame > 0: full_grid[0] += start_frame - # Shard the grid if using sequence parallelism (sp_world_size > 1) assert shard_dim < len( rope_dim_list), f"shard_dim {shard_dim} must be less than number of dimensions {len(rope_dim_list)}" if sp_world_size > 1: - # Get the shape of the full grid grid_shape = list(full_grid.shape[1:]) - # Ensure the dimension to shard is divisible by sp_world_size assert grid_shape[shard_dim] % sp_world_size == 0, ( f"Dimension {shard_dim} with size {grid_shape[shard_dim]} is not divisible " f"by sequence parallel world size {sp_world_size}") - # Compute the start and end indices for this rank's shard shard_size = grid_shape[shard_dim] // sp_world_size start_idx = sp_rank * shard_size end_idx = (sp_rank + 1) * shard_size - # Create slicing indices for each dimension slice_indices = [slice(None) for _ in range(len(grid_shape))] slice_indices[shard_dim] = slice(start_idx, end_idx) - # Shard the grid - # Update grid shape for the sharded dimension grid_shape[shard_dim] = grid_shape[shard_dim] // sp_world_size grid = torch.empty((len(rope_dim_list), ) + tuple(grid_shape), dtype=full_grid.dtype) for i in range(len(rope_dim_list)): @@ -431,7 +419,6 @@ def get_nd_rotary_pos_embed( assert len(interpolation_factor) == len( rope_dim_list), "len(interpolation_factor) should equal to len(rope_dim_list)" - # use 1/ndim of dimensions to encode grid_axis embs = [] for i in range(len(rope_dim_list)): emb = get_1d_rotary_pos_embed( @@ -506,7 +493,6 @@ def get_rotary_pos_embed( assert sum(rope_dim_list) == head_dim, "sum(rope_dim_list) should equal to head_dim of attention layer" - # Get SP info if do_sp_sharding: sp_group = get_sp_group() sp_rank = sp_group.rank_in_group diff --git a/fastvideo/layers/rotary_embedding_3d.py b/fastvideo/layers/rotary_embedding_3d.py index f284cf660..43dfc045e 100644 --- a/fastvideo/layers/rotary_embedding_3d.py +++ b/fastvideo/layers/rotary_embedding_3d.py @@ -99,17 +99,14 @@ class RotaryPositionalEmbedding3D(nn.Module): dim_h = 2 * (self.head_dim // 6) dim_w = 2 * (self.head_dim // 6) - # Compute frequency bands for each dimension freqs_t = 1.0 / (self.base**(torch.arange(0, dim_t, 2)[:(dim_t // 2)].float() / dim_t)) freqs_h = 1.0 / (self.base**(torch.arange(0, dim_h, 2)[:(dim_h // 2)].float() / dim_h)) freqs_w = 1.0 / (self.base**(torch.arange(0, dim_w, 2)[:(dim_w // 2)].float() / dim_w)) - # Create position grids grid_t = torch.arange(num_frames, dtype=torch.float32) grid_h = torch.arange(height, dtype=torch.float32) grid_w = torch.arange(width, dtype=torch.float32) - # Compute frequencies for each position freqs_t = torch.einsum("..., f -> ... f", grid_t, freqs_t) freqs_h = torch.einsum("..., f -> ... f", grid_h, freqs_h) freqs_w = torch.einsum("..., f -> ... f", grid_w, freqs_w) @@ -151,18 +148,15 @@ class RotaryPositionalEmbedding3D(nn.Module): Returns: (q_rotated, k_rotated): Rotated query and key tensors """ - # Register grid size if not cached if grid_size not in self.freqs_dict: self.register_grid_size(grid_size) - # Get cached frequencies freqs_cis = self.freqs_dict[grid_size].to(q.device) # Cast to float32 for precision q_, k_ = q.float(), k.float() freqs_cis = freqs_cis.float() - # Compute cos and sin cos = freqs_cis.cos() sin = freqs_cis.sin() @@ -170,7 +164,6 @@ class RotaryPositionalEmbedding3D(nn.Module): cos = rearrange(cos, "n d -> 1 1 n d") sin = rearrange(sin, "n d -> 1 1 n d") - # Apply rotation q_ = (q_ * cos) + (rotate_half(q_) * sin) k_ = (k_ * cos) + (rotate_half(k_) * sin) diff --git a/fastvideo/layers/utils.py b/fastvideo/layers/utils.py index b0d25d17d..eaca5e082 100644 --- a/fastvideo/layers/utils.py +++ b/fastvideo/layers/utils.py @@ -10,7 +10,6 @@ def get_token_bin_counts_and_mask( vocab_size: int, num_seqs: int, ) -> tuple[torch.Tensor, torch.Tensor]: - # Compute the bin counts for the tokens. # vocab_size + 1 for padding. bin_counts = torch.zeros((num_seqs, vocab_size + 1), dtype=torch.long, device=tokens.device) bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens)) diff --git a/fastvideo/layers/visual_embedding.py b/fastvideo/layers/visual_embedding.py index f544e0db0..b2ef9439d 100644 --- a/fastvideo/layers/visual_embedding.py +++ b/fastvideo/layers/visual_embedding.py @@ -34,7 +34,6 @@ class PatchEmbed(nn.Module): dtype=None, prefix: str = ""): super().__init__() - # Convert patch_size to 2-tuple if isinstance(patch_size, list | tuple): if len(patch_size) == 1: patch_size = (patch_size[0], patch_size[0]) @@ -129,7 +128,6 @@ class TimestepEmbedder(nn.Module): dtype=self.freq_dtype).to(self.mlp.fc_in.weight.dtype) if timestep_seq_len is not None: t_freq = t_freq.unflatten(0, (1, timestep_seq_len)) - # t_freq = t_freq.to(self.mlp.fc_in.weight.dtype) t_emb = self.mlp(t_freq) return t_emb @@ -240,7 +238,6 @@ def get_timestep_embedding( emb = torch.exp(exponent) emb = timesteps[:, None].float() * emb[None, :] - # scale embeddings emb = scale * emb # concat sine and cosine embeddings diff --git a/fastvideo/layers/vocab_parallel_embedding.py b/fastvideo/layers/vocab_parallel_embedding.py index 6f49fb311..63aaa0958 100644 --- a/fastvideo/layers/vocab_parallel_embedding.py +++ b/fastvideo/layers/vocab_parallel_embedding.py @@ -246,7 +246,6 @@ class VocabParallelEmbedding(torch.nn.Module): org_vocab_size_padded, tp_rank, tp_size)) padded_added_vocab_start_index, padded_added_vocab_end_index = (vocab_range_from_global_vocab_size( num_added_embeddings_padded, tp_rank, tp_size, offset=org_vocab_size)) - # remove padding org_vocab_start_index = min(padded_org_vocab_start_index, org_vocab_size) org_vocab_end_index = min(padded_org_vocab_end_index, org_vocab_size) added_vocab_start_index = min(padded_added_vocab_start_index, vocab_size) @@ -316,7 +315,6 @@ class VocabParallelEmbedding(torch.nn.Module): param.data.copy_(loaded_weight) return - # Shard indexes for loading the weight start_idx = self.shard_indices.org_vocab_start_index shard_size = self.shard_indices.org_vocab_end_index - start_idx @@ -330,7 +328,6 @@ class VocabParallelEmbedding(torch.nn.Module): else: assert loaded_weight.shape[output_dim] == self.org_vocab_size - # Copy the data. Select chunk corresponding to current shard. loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size) param[:loaded_weight.shape[0]].data.copy_(loaded_weight) @@ -338,7 +335,6 @@ class VocabParallelEmbedding(torch.nn.Module): def forward(self, input_): if self.tp_size > 1: - # Build the mask. masked_input, input_mask = get_masked_input_and_mask(input_, self.shard_indices.org_vocab_start_index, self.shard_indices.org_vocab_end_index, self.shard_indices.num_org_vocab_padding, @@ -346,7 +342,6 @@ class VocabParallelEmbedding(torch.nn.Module): self.shard_indices.added_vocab_end_index) else: masked_input = input_ - # Get the embeddings. output_parallel = self.quant_method.embedding(self, masked_input.long()) # Mask the output embedding. if self.tp_size > 1: diff --git a/fastvideo/logger.py b/fastvideo/logger.py index f00081ac1..8e9140fa7 100644 --- a/fastvideo/logger.py +++ b/fastvideo/logger.py @@ -70,7 +70,6 @@ def _print_info_once(logger: Logger, msg: str) -> None: @lru_cache def _print_warning_once(logger: Logger, msg: str) -> None: - # Set the stacklevel to 2 to print the original caller's line info logger.warning(msg, stacklevel=2) @@ -248,7 +247,6 @@ logger = init_logger(__name__) def _trace_calls(log_path, root_dir, frame, event, arg=None): if event in ['call', 'return']: - # Extract the filename, line number, function name, and the code object filename = frame.f_code.co_filename lineno = frame.f_lineno func_name = frame.f_code.co_name diff --git a/fastvideo/mlx_runtime/memory.py b/fastvideo/mlx_runtime/memory.py index 416d7cf8e..431babdff 100644 --- a/fastvideo/mlx_runtime/memory.py +++ b/fastvideo/mlx_runtime/memory.py @@ -168,7 +168,6 @@ def apply_memory_limits( mx_module = mx - # Apply each limit independently; record failures without stopping. limits = [ ("mlx_memory_limit", memory_bytes, mx_module.set_memory_limit), ("mlx_cache_limit", cache_bytes, mx_module.set_cache_limit), diff --git a/fastvideo/mlx_runtime/minimax_h3.py b/fastvideo/mlx_runtime/minimax_h3.py index bf4d66866..7c5b2e3b1 100644 --- a/fastvideo/mlx_runtime/minimax_h3.py +++ b/fastvideo/mlx_runtime/minimax_h3.py @@ -95,9 +95,7 @@ def linear(x, weight, bias=None): return _shared_linear(x, weight, bias, use_affine_dq_gemm=True) -# --------------------------------------------------------------------------- # Constants (mirrors fastvideo/pipelines/basic/minimax_h3/packing.py) -# --------------------------------------------------------------------------- MINIMAX_H3_VIDEO_TAG = 0 MINIMAX_H3_TEXT_TAG = 1 @@ -165,9 +163,7 @@ def _is_ignored_dense_key(key: str, *, include_vsa: bool = False) -> bool: return any(part in key for part in _IGNORED_DENSE_KEY_PARTS) -# --------------------------------------------------------------------------- # Geometry (numpy, float64 — the reference RoPE coordinate arithmetic) -# --------------------------------------------------------------------------- def align_num_frames(num_frames: int) -> int: @@ -618,9 +614,7 @@ def build_row_timesteps( return unique.astype(np.float32), inverse.astype(np.int64) -# --------------------------------------------------------------------------- # Scheduler (rectified flow, clean-time convention, data-ward velocity) -# --------------------------------------------------------------------------- def _unique_consecutive(values: np.ndarray) -> np.ndarray: @@ -943,9 +937,6 @@ class MiniMaxH3SchedulerState: return timestep * sample + (1.0 - timestep) * noise -# --------------------------------------------------------------------------- -# DiT -# --------------------------------------------------------------------------- def _as_mx_indices(indices): @@ -1737,9 +1728,6 @@ class MLXMiniMaxH3DiT: __call__ = forward -# --------------------------------------------------------------------------- -# Checkpoint loading (diffusers-layout safetensors) and pre-quantized save/load -# --------------------------------------------------------------------------- def _load_array(array, dtype): diff --git a/fastvideo/mlx_runtime/minimax_h3_audio_vae.py b/fastvideo/mlx_runtime/minimax_h3_audio_vae.py index a333395b6..771e7f00e 100644 --- a/fastvideo/mlx_runtime/minimax_h3_audio_vae.py +++ b/fastvideo/mlx_runtime/minimax_h3_audio_vae.py @@ -88,9 +88,7 @@ class MiniMaxH3AudioVAEConfigView: ACTIVATION_RATIO = 2 -# --------------------------------------------------------------------------- # Primitives (torch layout (B, C, L); convs run channels-last internally) -# --------------------------------------------------------------------------- def _wn_weight(weight_v, weight_g): @@ -237,9 +235,6 @@ def _activation1d(x, alpha, beta, up_filter, down_filter, ratio: int = ACTIVATIO return _low_pass(x, down_filter, ratio) -# --------------------------------------------------------------------------- -# The MLX H3 audio VAE -# --------------------------------------------------------------------------- class MLXMiniMaxH3AudioVAE: diff --git a/fastvideo/models/audio/ltx2_audio_vae.py b/fastvideo/models/audio/ltx2_audio_vae.py index 56df5374c..1ac1557b9 100644 --- a/fastvideo/models/audio/ltx2_audio_vae.py +++ b/fastvideo/models/audio/ltx2_audio_vae.py @@ -14,9 +14,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -# ============================================================================= -# Enums -# ============================================================================= class NormType(Enum): @@ -40,9 +37,6 @@ class CausalityAxis(Enum): WIDTH_COMPATIBILITY = "width-compatibility" -# ============================================================================= -# Audio Latent Shape -# ============================================================================= class AudioLatentShape(NamedTuple): @@ -65,16 +59,10 @@ class AudioLatentShape(NamedTuple): ) -# ============================================================================= -# Constants -# ============================================================================= LATENT_DOWNSAMPLE_FACTOR = 4 LRELU_SLOPE = 0.1 -# ============================================================================= -# Normalization Layers -# ============================================================================= class PixelNorm(nn.Module): @@ -107,9 +95,6 @@ def build_normalization_layer( raise ValueError(f"Invalid normalization type: {normtype}") -# ============================================================================= -# Per-Channel Statistics -# ============================================================================= class PerChannelStatistics(nn.Module): @@ -129,9 +114,6 @@ class PerChannelStatistics(nn.Module): return (x - self.get_buffer("mean-of-means").to(x)) / self.get_buffer("std-of-means").to(x) -# ============================================================================= -# Audio Patchifier -# ============================================================================= class AudioPatchifier: @@ -165,9 +147,6 @@ class AudioPatchifier: ) -# ============================================================================= -# Causal 2D Convolution -# ============================================================================= class CausalConv2d(nn.Module): @@ -191,11 +170,9 @@ class CausalConv2d(nn.Module): self.causality_axis = causality_axis - # Ensure kernel_size and dilation are tuples kernel_size = nn.modules.utils._pair(kernel_size) dilation = nn.modules.utils._pair(dilation) - # Calculate padding dimensions pad_h = (kernel_size[0] - 1) * dilation[0] pad_w = (kernel_size[1] - 1) * dilation[1] @@ -255,9 +232,6 @@ def make_conv2d( ) -# ============================================================================= -# Attention Blocks -# ============================================================================= class AttnBlock(nn.Module): @@ -284,7 +258,6 @@ class AttnBlock(nn.Module): k = self.k(h_) v = self.v(h_) - # Compute attention b, c, h, w = q.shape q = q.reshape(b, c, h * w).contiguous() q = q.permute(0, 2, 1).contiguous() # b, hw, c @@ -320,9 +293,6 @@ def make_attn( raise ValueError(f"Unknown attention type: {attn_type}") -# ============================================================================= -# ResNet Blocks (2D for Audio VAE) -# ============================================================================= class ResnetBlock(nn.Module): @@ -396,9 +366,6 @@ class ResnetBlock(nn.Module): return x + h -# ============================================================================= -# Vocoder ResBlocks (1D) -# ============================================================================= class ResBlock1(nn.Module): @@ -455,9 +422,6 @@ class ResBlock2(nn.Module): return x -# ============================================================================= -# BigVGAN v2 / BWE helpers -# ============================================================================= def get_padding(kernel_size: int, dilation: int = 1) -> int: @@ -737,9 +701,6 @@ class AMPBlock1(nn.Module): return x -# ============================================================================= -# Downsampling -# ============================================================================= class Downsample(nn.Module): @@ -835,9 +796,6 @@ def build_downsampling_path( return down_modules, block_in -# ============================================================================= -# Upsampling -# ============================================================================= class Upsample(nn.Module): @@ -924,9 +882,6 @@ def build_upsampling_path( return up_modules, block_in -# ============================================================================= -# Mid Block -# ============================================================================= def build_mid_block( @@ -967,9 +922,6 @@ def run_mid_block(mid: nn.Module, features: torch.Tensor) -> torch.Tensor: return mid.block_2(features, temb=None) -# ============================================================================= -# Audio Encoder -# ============================================================================= class AudioEncoder(nn.Module): @@ -1125,9 +1077,6 @@ class AudioEncoder(nn.Module): return self.patchifier.unpatchify(latent_normalized, latent_shape) -# ============================================================================= -# Audio Decoder -# ============================================================================= class AudioDecoder(nn.Module): @@ -1290,11 +1239,9 @@ class AudioDecoder(nn.Module): decoded_output = decoded_output[:, :target_channels, :min(current_time, target_time ), :min(current_freq, target_freq)] - # Calculate padding needed time_padding_needed = target_time - decoded_output.shape[2] freq_padding_needed = target_freq - decoded_output.shape[3] - # Apply padding if needed if time_padding_needed > 0 or freq_padding_needed > 0: padding = ( 0, @@ -1332,9 +1279,6 @@ class AudioDecoder(nn.Module): return torch.tanh(h) if self.tanh_out else h -# ============================================================================= -# Vocoder -# ============================================================================= class Vocoder(nn.Module): @@ -1359,7 +1303,6 @@ class Vocoder(nn.Module): ): super().__init__() - # Initialize default values if resblock_kernel_sizes is None: resblock_kernel_sizes = [3, 7, 11] if upsample_rates is None: @@ -1459,9 +1402,6 @@ class Vocoder(nn.Module): return x -# ============================================================================= -# BWE STFT / MelSTFT / VocoderWithBWE -# ============================================================================= def _build_stft_basis(filter_length: int, win_length: int) -> torch.Tensor: @@ -1645,9 +1585,6 @@ class VocoderWithBWE(nn.Module): return torch.clamp(residual + skip, -1, 1)[..., :output_length] -# ============================================================================= -# Configurators -# ============================================================================= class AudioEncoderConfigurator: @@ -1804,9 +1741,6 @@ class VocoderConfigurator: ) -# ============================================================================= -# Public API Wrappers -# ============================================================================= class LTX2AudioEncoder(nn.Module): diff --git a/fastvideo/models/dits/base.py b/fastvideo/models/dits/base.py index 5a4103e9d..471700910 100644 --- a/fastvideo/models/dits/base.py +++ b/fastvideo/models/dits/base.py @@ -9,7 +9,6 @@ from fastvideo.configs.models import DiTConfig from fastvideo.platforms import AttentionBackendEnum -# TODO class BaseDiT(nn.Module, ABC): _fsdp_shard_conditions: list = [] _compile_conditions: list = [] diff --git a/fastvideo/models/dits/cosmos.py b/fastvideo/models/dits/cosmos.py index 1ef62ba8a..ac3a8630c 100644 --- a/fastvideo/models/dits/cosmos.py +++ b/fastvideo/models/dits/cosmos.py @@ -171,7 +171,6 @@ class CosmosSelfAttention(nn.Module): if encoder_hidden_states is None: encoder_hidden_states = hidden_states - # Get QKV query = self.to_q(hidden_states) key = self.to_k(encoder_hidden_states) value = self.to_v(encoder_hidden_states) @@ -181,18 +180,15 @@ class CosmosSelfAttention(nn.Module): key = key.unflatten(2, (self.num_heads, -1)).transpose(1, 2) value = value.unflatten(2, (self.num_heads, -1)).transpose(1, 2) - # Apply normalization if self.norm_q is not None: query = self.norm_q(query) if self.norm_k is not None: key = self.norm_k(key) - # Apply RoPE if provided if image_rotary_emb is not None: query = apply_rotary_emb(query, image_rotary_emb, use_real=True, use_real_unbind_dim=-2) key = apply_rotary_emb(key, image_rotary_emb, use_real=True, use_real_unbind_dim=-2) - # Prepare for GQA (Grouped Query Attention) if torch.onnx.is_in_onnx_export(): query_idx = torch.tensor(query.size(3), device=query.device) key_idx = torch.tensor(key.size(3), device=key.device) @@ -252,7 +248,6 @@ class CosmosCrossAttention(nn.Module): encoder_hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor: - # Get QKV query = self.to_q(hidden_states) key = self.to_k(encoder_hidden_states) value = self.to_v(encoder_hidden_states) @@ -262,13 +257,11 @@ class CosmosCrossAttention(nn.Module): key = key.unflatten(2, (self.num_heads, -1)).transpose(1, 2) value = value.unflatten(2, (self.num_heads, -1)).transpose(1, 2) - # Apply normalization if self.norm_q is not None: query = self.norm_q(query) if self.norm_k is not None: key = self.norm_k(key) - # Prepare for GQA (Grouped Query Attention) if torch.onnx.is_in_onnx_export(): query_idx = torch.tensor(query.size(3), device=query.device) key_idx = torch.tensor(key.size(3), device=key.device) @@ -462,7 +455,6 @@ class CosmosLearnablePositionalEmbed(nn.Module): class CosmosTransformer3DModel(BaseDiT): _fsdp_shard_conditions = CosmosVideoConfig()._fsdp_shard_conditions _compile_conditions = CosmosVideoConfig()._compile_conditions - # _supported_attention_backends = CosmosVideoConfig()._supported_attention_backends param_names_mapping = CosmosVideoConfig().param_names_mapping lora_param_names_mapping = CosmosVideoConfig().lora_param_names_mapping diff --git a/fastvideo/models/dits/cosmos2_5.py b/fastvideo/models/dits/cosmos2_5.py index c35bd88e9..c722eeed5 100644 --- a/fastvideo/models/dits/cosmos2_5.py +++ b/fastvideo/models/dits/cosmos2_5.py @@ -184,7 +184,6 @@ class Cosmos25AdaLayerNormZero(nn.Module): Returns: normalized_hidden_states: Modulated normalized hidden states """ - # Apply layer norm and modulation hidden_states = self.norm(hidden_states) hidden_states = hidden_states * (1 + scale) + shift return hidden_states @@ -239,7 +238,6 @@ class Cosmos25SelfAttention(nn.Module): hidden_states: (B, S, D) where S = T*H*W rope_emb: Tuple of (cos, sin) for RoPE """ - # Get QKV query, _ = self.to_q(hidden_states) key, _ = self.to_k(hidden_states) value, _ = self.to_v(hidden_states) @@ -249,11 +247,9 @@ class Cosmos25SelfAttention(nn.Module): key = key.unflatten(-1, (self.num_heads, self.head_dim)).transpose(1, 2) value = value.unflatten(-1, (self.num_heads, self.head_dim)).transpose(1, 2) - # Apply QK normalization query = self.norm_q(query) key = self.norm_k(key) - # Apply RoPE if provided (query/key are now in (B, H, S, D_h) format) if rope_emb is not None: cos, sin = rope_emb query = apply_rotary_emb(query, (cos, sin), use_real=True, use_real_unbind_dim=-2) @@ -326,7 +322,6 @@ class Cosmos25CrossAttention(nn.Module): hidden_states: (B, S, D) encoder_hidden_states: (B, N, D_text) """ - # Get QKV query, _ = self.to_q(hidden_states) key, _ = self.to_k(encoder_hidden_states) value, _ = self.to_v(encoder_hidden_states) @@ -336,7 +331,6 @@ class Cosmos25CrossAttention(nn.Module): key = key.unflatten(-1, (self.num_heads, self.head_dim)) value = value.unflatten(-1, (self.num_heads, self.head_dim)) - # Apply QK normalization query = self.norm_q(query) key = self.norm_k(key) @@ -439,13 +433,11 @@ class Cosmos25TransformerBlock(nn.Module): rope_emb: Tuple of (cos, sin) for RoPE extra_pos_emb: Optional learnable positional embeddings """ - # Add extra positional embeddings if provided if extra_pos_emb is not None: hidden_states = hidden_states + extra_pos_emb B, T, H, W, D = hidden_states.shape - # Step 1: Compute ALL modulation parameters once (matches official model) if self.use_adaln_lora and adaln_lora is not None: shift_self_attn, scale_self_attn, gate_self_attn = (self.adaln_modulation_self_attn(embedded_timestep) + adaln_lora).chunk(3, dim=-1) @@ -473,7 +465,6 @@ class Cosmos25TransformerBlock(nn.Module): scale_mlp = scale_mlp.unsqueeze(2).unsqueeze(2).type_as(hidden_states) gate_mlp = gate_mlp.unsqueeze(2).unsqueeze(2).type_as(hidden_states) - # Step 2: Self-attention block norm_hidden_states = self.norm1(hidden_states, shift_self_attn, scale_self_attn) # Flatten for attention: (B, T, H, W, D) -> (B, THW, D) norm_hidden_states_flat = norm_hidden_states.flatten(1, 3) @@ -484,7 +475,6 @@ class Cosmos25TransformerBlock(nn.Module): attn_output = attn_output.unflatten(1, (T, H, W)) # (B, T, H, W, D) hidden_states = hidden_states + gate_self_attn * attn_output - # Step 3: Cross-attention block norm_hidden_states = self.norm2(hidden_states, shift_cross_attn, scale_cross_attn) norm_hidden_states_flat = norm_hidden_states.flatten(1, 3) @@ -497,7 +487,6 @@ class Cosmos25TransformerBlock(nn.Module): attn_output = attn_output.unflatten(1, (T, H, W)) hidden_states = hidden_states + gate_cross_attn * attn_output - # Step 4: MLP block norm_hidden_states = self.norm3(hidden_states, shift_mlp, scale_mlp) norm_hidden_states_flat = norm_hidden_states.flatten(1, 3) @@ -557,14 +546,12 @@ class Cosmos25RotaryPosEmbed(nn.Module): # T, H, W are already patch dimensions after patch_embed # No need to divide by patch_size - # Generate frequency scales with NTK h_theta = 10000.0 * self.h_ntk_factor w_theta = 10000.0 * self.w_ntk_factor t_theta = 10000.0 * self.t_ntk_factor seq = torch.arange(max(self.max_size), device=device, dtype=torch.float32) - # Use self.dim_h/w/t which were set during initialization dim_h_range = torch.arange(0, self.dim_h, 2, device=device, dtype=torch.float32)[:(self.dim_h // 2)] / self.dim_h dim_w_range = torch.arange(0, self.dim_w, 2, device=device, @@ -576,12 +563,10 @@ class Cosmos25RotaryPosEmbed(nn.Module): w_spatial_freqs = 1.0 / (w_theta**dim_w_range) temporal_freqs = 1.0 / (t_theta**dim_t_range) - # Generate positional embeddings half_emb_h = torch.outer(seq[:H], h_spatial_freqs) half_emb_w = torch.outer(seq[:W], w_spatial_freqs) if self.enable_fps_modulation and fps is not None: - # Apply FPS scaling half_emb_t = torch.outer(seq[:T] / fps * self.base_fps, temporal_freqs) else: half_emb_t = torch.outer(seq[:T], temporal_freqs) @@ -638,7 +623,6 @@ class Cosmos25LearnablePositionalEmbed(nn.Module): emb = emb_t + emb_h + emb_w - # Normalize norm = torch.linalg.vector_norm(emb, dim=-1, keepdim=True, dtype=torch.float32) norm = torch.add(self.eps, norm, alpha=np.sqrt(norm.numel() / emb.numel())) return (emb / norm).type_as(hidden_states) @@ -687,18 +671,15 @@ class Cosmos25FinalLayer(nn.Module): embedded_timestep: (B, T, D) or (B, D) adaln_lora: (B, T, 3D) or None """ - # Generate modulation parameters embedded_timestep = self.activation(embedded_timestep) embedded_timestep = self.linear_1(embedded_timestep) embedded_timestep = self.linear_2(embedded_timestep) if self.use_adaln_lora and adaln_lora is not None: - # Use first 2*hidden_size elements for shift/scale embedded_timestep = embedded_timestep + adaln_lora[..., :2 * self.hidden_size] shift, scale = embedded_timestep.chunk(2, dim=-1) - # Apply normalization and modulation hidden_states = self.norm(hidden_states) # Reshape for broadcasting if needed @@ -888,14 +869,12 @@ class Cosmos25Transformer3DModel(BaseDiT): else: raise ValueError(f"Unsupported timestep shape: {timestep.shape}") - # Now timestep is always (B, T), pass directly to time_embed embedded_timestep, adaln_lora = self.time_embed(hidden_states, timestep) # 7. Apply cross-attention projection (if used) if self.use_crossattn_projection: encoder_hidden_states = self.crossattn_proj(encoder_hidden_states.to(self.crossattn_proj[0].weight.dtype)) - # Prepare attention mask if attention_mask is not None: attention_mask = attention_mask.unsqueeze(1).unsqueeze(1) # (B, 1, 1, N) diff --git a/fastvideo/models/dits/flux_2.py b/fastvideo/models/dits/flux_2.py index de1ea4cb4..fe75619f9 100644 --- a/fastvideo/models/dits/flux_2.py +++ b/fastvideo/models/dits/flux_2.py @@ -420,7 +420,6 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): dim=-1, ) - # Handle the attention logic query, key, value = qkv.chunk(3, dim=-1) query = query.unflatten(-1, (self.heads, -1)) @@ -441,7 +440,6 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) - # Handle the feedforward (FF) logic mlp_hidden_states = self.mlp_act_fn(mlp_hidden_states) # Concatenate and parallel output projection @@ -605,7 +603,6 @@ class Flux2TransformerBlock(nn.Module): attn_output, context_attn_output = attention_outputs - # Process attention outputs for the image stream (`hidden_states`). attn_output = gate_msa * attn_output hidden_states = hidden_states + attn_output @@ -615,7 +612,6 @@ class Flux2TransformerBlock(nn.Module): ff_output = self.ff(norm_hidden_states) hidden_states = hidden_states + gate_mlp * ff_output - # Process attention outputs for the text stream (`encoder_hidden_states`). context_attn_output = c_gate_msa * context_attn_output encoder_hidden_states = encoder_hidden_states + context_attn_output @@ -685,7 +681,6 @@ class Flux2Modulation(nn.Module): if mod.ndim == 2: mod = mod.unsqueeze(1) mod_params = torch.chunk(mod, 3 * self.mod_param_sets, dim=-1) - # Return tuple of 3-tuples of modulation params shift/scale/gate return tuple(mod_params[3 * i:3 * (i + 1)] for i in range(self.mod_param_sets)) @@ -815,7 +810,6 @@ class Flux2Transformer2DModel(BaseDiT): self.inner_dim = num_attention_heads * attention_head_dim self.guidance_embeds = guidance_embeds - # Add required BaseDiT attributes self.hidden_size = self.inner_dim self.num_attention_heads = num_attention_heads self.num_channels_latents = self.out_channels @@ -1006,7 +1000,6 @@ class Flux2Transformer2DModel(BaseDiT): freqs_cis=freqs_cis, joint_attention_kwargs=joint_attention_kwargs, ) - # Remove text tokens from concatenated stream hidden_states = hidden_states[:, num_txt_tokens:, ...] # 6. Output layers diff --git a/fastvideo/models/dits/gen3c.py b/fastvideo/models/dits/gen3c.py index e00cf9795..1c138a2fd 100644 --- a/fastvideo/models/dits/gen3c.py +++ b/fastvideo/models/dits/gen3c.py @@ -213,7 +213,6 @@ class Gen3CAdaLayerNormZero(nn.Module): Returns: normalized_hidden_states: Modulated normalized hidden states """ - # Apply layer norm and modulation hidden_states = self.norm(hidden_states) hidden_states = hidden_states * (1 + scale) + shift return hidden_states @@ -267,7 +266,6 @@ class Gen3CSelfAttention(nn.Module): rope_emb: Tuple of (cos, sin) for RoPE original_seq_len: Original (unpadded) full sequence length """ - # Get QKV query, _ = self.to_q(hidden_states) key, _ = self.to_k(hidden_states) value, _ = self.to_v(hidden_states) @@ -277,11 +275,9 @@ class Gen3CSelfAttention(nn.Module): key = key.unflatten(-1, (self.num_heads, self.head_dim)).transpose(1, 2) value = value.unflatten(-1, (self.num_heads, self.head_dim)).transpose(1, 2) - # Apply QK normalization query = self.norm_q(query) key = self.norm_k(key) - # Apply RoPE if provided (query/key are now in (B, H, S, D_h) format) if rope_emb is not None: cos, sin = rope_emb query = apply_rotary_emb(query, (cos, sin), use_real=True, use_real_unbind_dim=-2) @@ -350,7 +346,6 @@ class Gen3CCrossAttention(nn.Module): hidden_states: (B, S, D) encoder_hidden_states: (B, N, D_text) """ - # Get QKV query, _ = self.to_q(hidden_states) key, _ = self.to_k(encoder_hidden_states) value, _ = self.to_v(encoder_hidden_states) @@ -360,7 +355,6 @@ class Gen3CCrossAttention(nn.Module): key = key.unflatten(-1, (self.num_heads, self.head_dim)) value = value.unflatten(-1, (self.num_heads, self.head_dim)) - # Apply QK normalization query = self.norm_q(query) key = self.norm_k(key) @@ -460,13 +454,11 @@ class Gen3CTransformerBlock(nn.Module): extra_pos_emb: Optional learnable positional embeddings (B, S, D) original_seq_len: Original (unpadded) full sequence length """ - # Add extra positional embeddings if provided if extra_pos_emb is not None: hidden_states = hidden_states + extra_pos_emb B, S, D = hidden_states.shape - # Compute modulation parameters if self.use_adaln_lora and adaln_lora is not None: shift_self_attn, scale_self_attn, gate_self_attn = ( _run_linear_stack(self.adaln_modulation_self_attn, affine_emb) + adaln_lora).chunk(3, dim=-1) @@ -560,7 +552,6 @@ class Gen3CRotaryPosEmbed(nn.Module): batch_size, T, H, W, input_dim = hidden_states.shape device = hidden_states.device - # Generate frequency scales with NTK h_theta = 10000.0 * self.h_ntk_factor w_theta = 10000.0 * self.w_ntk_factor t_theta = 10000.0 * self.t_ntk_factor @@ -578,7 +569,6 @@ class Gen3CRotaryPosEmbed(nn.Module): w_spatial_freqs = 1.0 / (w_theta**dim_w_range) temporal_freqs = 1.0 / (t_theta**dim_t_range) - # Generate positional embeddings half_emb_h = torch.outer(seq[:H], h_spatial_freqs) half_emb_w = torch.outer(seq[:W], w_spatial_freqs) @@ -639,7 +629,6 @@ class Gen3CLearnablePositionalEmbed(nn.Module): emb = emb_t + emb_h + emb_w - # Normalize norm = torch.linalg.vector_norm(emb, dim=-1, keepdim=True, dtype=torch.float32) norm = torch.add(self.eps, norm, alpha=np.sqrt(norm.numel() / emb.numel())) return (emb / norm).type_as(hidden_states) @@ -693,7 +682,6 @@ class Gen3CFinalLayer(nn.Module): affine_emb: (B, D) adaln_lora: (B, 3D) or None """ - # Generate modulation parameters modulation = _run_linear_stack(self.adaln_modulation, affine_emb) if self.use_adaln_lora and adaln_lora is not None: @@ -701,7 +689,6 @@ class Gen3CFinalLayer(nn.Module): shift, scale = modulation.chunk(2, dim=-1) - # Apply normalization and modulation hidden_states = self.norm(hidden_states) # Broadcast over all non-channel dimensions. Supports both: diff --git a/fastvideo/models/dits/hunyuanvideo.py b/fastvideo/models/dits/hunyuanvideo.py index e69d56a7a..620769c26 100644 --- a/fastvideo/models/dits/hunyuanvideo.py +++ b/fastvideo/models/dits/hunyuanvideo.py @@ -9,7 +9,6 @@ from fastvideo.attention import DistributedAttention, LocalAttention from fastvideo.configs.models.dits import HunyuanVideoConfig from fastvideo.layers.layernorm import (LayerNormScaleShift, ScaleResidual, ScaleResidualLayerNormScaleShift) from fastvideo.layers.linear import ReplicatedLinear -# TODO(will-PY-refactor): RMSNorm .... from fastvideo.layers.mlp import MLP from fastvideo.layers.rotary_embedding import get_rotary_pos_embed from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, TimestepEmbedder, unpatchify) @@ -198,9 +197,7 @@ class MMDoubleStreamBlock(nn.Module): txt_mlp_gate, ) = torch.chunk(txt_mod_outputs, 6, dim=-1) - # Prepare image for attention using fused operation img_attn_input = self.img_attn_norm(img, img_attn_shift, img_attn_scale) - # Get QKV for image img_qkv, _ = self.img_attn_qkv(img_attn_input) batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1] @@ -208,15 +205,12 @@ class MMDoubleStreamBlock(nn.Module): img_qkv = img_qkv.view(batch_size, image_seq_len, 3, self.num_attention_heads, -1) img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, 2] - # Apply QK-Norm if needed img_q = self.img_attn_q_norm(img_q).to(img_v) img_k = self.img_attn_k_norm(img_k).to(img_v) - # Prepare text for attention using fused operation txt_attn_input = self.txt_attn_norm(txt, txt_attn_shift, txt_attn_scale) - # Get QKV for text txt_qkv, _ = self.txt_attn_qkv(txt_attn_input) batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1] @@ -224,29 +218,22 @@ class MMDoubleStreamBlock(nn.Module): txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3, self.num_attention_heads, -1) txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, 2] - # Apply QK-Norm if needed txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype) txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype) - # Run distributed attention img_attn, txt_attn = self.attn(img_q, img_k, img_v, original_seq_len, txt_q, txt_k, txt_v, freqs_cis=freqs_cis) img_attn_out, _ = self.img_attn_proj(img_attn.reshape(batch_size, image_seq_len, -1)) - # Use fused operation for residual connection, normalization, and modulation img_mlp_input, img_residual = self.img_attn_residual_mlp_norm(img, img_attn_out, img_attn_gate, img_mlp_shift, img_mlp_scale) - # Process image MLP img_mlp_out = self.img_mlp(img_mlp_input) img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate) - # Process text attention output txt_attn_out, _ = self.txt_attn_proj(txt_attn.reshape(batch_size, text_seq_len, -1)) - # Use fused operation for residual connection, normalization, and modulation txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm(txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, txt_mlp_scale) - # Process text MLP txt_mlp_out = self.txt_mlp(txt_mlp_input) txt = self.txt_mlp_residual(txt_residual, txt_mlp_out, txt_mlp_gate) @@ -329,24 +316,19 @@ class MMSingleStreamBlock(nn.Module): freqs_cis: tuple[torch.Tensor, torch.Tensor], original_seq_len: int | None = None, ) -> torch.Tensor: - # Process modulation mod_shift, mod_scale, mod_gate = self.modulation(vec).unsqueeze(1).chunk(3, dim=-1) - # Apply pre-norm and modulation using fused operation x_mod = self.input_norm_scale_shift(x, mod_shift, mod_scale) - # Get combined projections linear1_out, _ = self.linear1(x_mod) # Split into QKV and MLP parts qkv, mlp = torch.split(linear1_out, [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1) - # Process QKV batch_size, seq_len = qkv.shape[0], qkv.shape[1] qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1) q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] - # Apply QK-Norm q = self.q_norm(q).to(v.dtype) k = self.k_norm(k).to(v.dtype) @@ -355,7 +337,6 @@ class MMSingleStreamBlock(nn.Module): img_k, txt_k = k[:, :-txt_len], k[:, -txt_len:] img_v, txt_v = v[:, :-txt_len], v[:, -txt_len:] - # Run distributed attention img_attn_output, txt_attn_output = self.attn(img_q, img_k, img_v, @@ -365,7 +346,6 @@ class MMSingleStreamBlock(nn.Module): txt_v, freqs_cis=freqs_cis) attn_output = torch.cat((img_attn_output, txt_attn_output), dim=1).view(batch_size, seq_len, -1) - # Process MLP activation mlp_output = self.mlp_act(mlp) # Combine attention and MLP outputs @@ -374,7 +354,6 @@ class MMSingleStreamBlock(nn.Module): # Final projection output, _ = self.linear2(combined) - # Apply residual connection with gating using fused operation return self.output_residual(x, output, mod_gate) @@ -391,7 +370,6 @@ class HunyuanVideoTransformer3DModel(BaseDiT): """ # PY: we make the input args the same as HF config - # shard single stream, double stream blocks, and refiner_blocks _fsdp_shard_conditions = HunyuanVideoConfig()._fsdp_shard_conditions _compile_conditions = HunyuanVideoConfig()._compile_conditions _supported_attention_backends = HunyuanVideoConfig()._supported_attention_backends @@ -412,7 +390,6 @@ class HunyuanVideoTransformer3DModel(BaseDiT): self.rope_theta = config.rope_theta self.text_states_dim = config.text_embed_dim self.text_states_dim_2 = config.pooled_projection_dim - # TODO(will): hack? self.dtype = config.dtype pe_dim = config.hidden_size // config.num_attention_heads @@ -485,8 +462,6 @@ class HunyuanVideoTransformer3DModel(BaseDiT): self.__post_init__() - # TODO: change the input the FORWAD_BACTCH Dict - # TODO: change output to a dict def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], @@ -521,7 +496,6 @@ class HunyuanVideoTransformer3DModel(BaseDiT): txt = encoder_hidden_states[0] text_states_2 = encoder_hidden_states[1] - # Get spatial dimensions _, _, ot, oh, ow = x.shape # codespell:ignore tt, th, tw = ( ot // self.patch_size[0], # codespell:ignore @@ -529,18 +503,14 @@ class HunyuanVideoTransformer3DModel(BaseDiT): ow // self.patch_size[2], ) - # Get rotary embeddings freqs_cos, freqs_sin = get_rotary_pos_embed((tt, th, tw), self.hidden_size, self.num_attention_heads, self.rope_dim_list, self.rope_theta) freqs_cos = freqs_cos.to(x.device) freqs_sin = freqs_sin.to(x.device) - # Prepare modulation vectors vec = self.time_in(t) - # Add text modulation vec = vec + self.vector_in(text_states_2) - # Add guidance modulation if needed if self.guidance_in and guidance is not None: vec = vec + self.guidance_in(guidance) # Embed image and text @@ -552,14 +522,12 @@ class HunyuanVideoTransformer3DModel(BaseDiT): freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None - # Process through double stream blocks for index, block in enumerate(self.double_blocks): double_block_args = [img, txt, vec, freqs_cis, original_seq_len] img, txt = block(*double_block_args) # Merge txt and img to pass through single stream blocks x = torch.cat((img, txt), 1) - # Process through single stream blocks if len(self.single_blocks) > 0: for index, block in enumerate(self.single_blocks): single_block_args = [ @@ -571,7 +539,6 @@ class HunyuanVideoTransformer3DModel(BaseDiT): ] x = block(*single_block_args) - # Extract image features img = x[:, :img_seq_len, ...] # Final layer processing @@ -631,10 +598,8 @@ class SingleTokenRefiner(nn.Module): ]) def forward(self, x, t): - # Get timestep embeddings timestep_aware_representations = self.t_embedder(t) - # Get context-aware representations context_aware_representations = torch.mean(x, dim=1) @@ -642,7 +607,6 @@ class SingleTokenRefiner(nn.Module): c = timestep_aware_representations + context_aware_representations # Project input x, _ = self.input_embedder(x) - # Process through refiner blocks for block in self.refiner_blocks: x = block(x, c) return x @@ -696,12 +660,10 @@ class IndividualTokenRefinerBlock(nn.Module): self.attn = LocalAttention( num_heads=num_attention_heads, head_size=hidden_size // num_attention_heads, - # TODO: remove hardcode supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), ) def forward(self, x, c): - # Get modulation parameters gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=-1) # Self-attention norm_x = self.norm1(x) @@ -711,7 +673,6 @@ class IndividualTokenRefinerBlock(nn.Module): qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1) q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] - # Run scaled dot product attention attn_output = self.attn(q, k, v) # [B, L, H, D] attn_output = attn_output.reshape(batch_size, seq_len, -1) # [B, L, H*D] diff --git a/fastvideo/models/dits/hunyuanvideo15.py b/fastvideo/models/dits/hunyuanvideo15.py index 4b0f08b4c..a2809c021 100644 --- a/fastvideo/models/dits/hunyuanvideo15.py +++ b/fastvideo/models/dits/hunyuanvideo15.py @@ -24,7 +24,6 @@ from fastvideo.distributed.communication_op import (sequence_model_parallel_all_ from fastvideo.configs.models.dits import HunyuanVideo15Config from fastvideo.layers.layernorm import (LayerNormScaleShift, ScaleResidual, ScaleResidualLayerNormScaleShift) from fastvideo.layers.linear import ReplicatedLinear -# TODO(will-PY-refactor): RMSNorm .... from fastvideo.layers.mlp import MLP from fastvideo.layers.rotary_embedding import get_rotary_pos_embed from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, TimestepEmbedder, unpatchify) @@ -293,9 +292,7 @@ class MMDoubleStreamBlock(nn.Module): txt_mlp_gate, ) = torch.chunk(txt_mod_outputs, 6, dim=-1) - # Prepare image for attention using fused operation img_attn_input = self.img_attn_norm(img, img_attn_shift, img_attn_scale) - # Get QKV for image img_qkv, _ = self.img_attn_qkv(img_attn_input) batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1] @@ -303,43 +300,34 @@ class MMDoubleStreamBlock(nn.Module): img_qkv = img_qkv.view(batch_size, image_seq_len, 3, self.num_attention_heads, -1) img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, 2] - # Apply QK-Norm if needed img_q = self.img_attn_q_norm(img_q).to(img_v) img_k = self.img_attn_k_norm(img_k).to(img_v) - # Prepare text for attention using fused operation txt_attn_input = self.txt_attn_norm(txt, txt_attn_shift, txt_attn_scale) - # Get QKV for text txt_qkv, _ = self.txt_attn_qkv(txt_attn_input) batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1] # Split QKV txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3, self.num_attention_heads, -1) txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, 2] - # Apply QK-Norm if needed txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype) txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype) img_attn, txt_attn = self.attn(img_q, img_k, img_v, original_seq_len, txt_q, txt_k, txt_v, freqs_cis=freqs_cis) img_attn_out, _ = self.img_attn_proj(img_attn.reshape(batch_size, image_seq_len, -1)) - # Use fused operation for residual connection, normalization, and modulation img_mlp_input, img_residual = self.img_attn_residual_mlp_norm(img, img_attn_out, img_attn_gate, img_mlp_shift, img_mlp_scale) - # Process image MLP img_mlp_out = self.img_mlp(img_mlp_input) img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate) - # Process text attention output txt_attn_out, _ = self.txt_attn_proj(txt_attn.reshape(batch_size, text_seq_len, -1)) - # Use fused operation for residual connection, normalization, and modulation txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm(txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, txt_mlp_scale) - # Process text MLP txt_mlp_out = self.txt_mlp(txt_mlp_input) txt = self.txt_mlp_residual(txt_residual, txt_mlp_out, txt_mlp_gate) @@ -351,7 +339,6 @@ class HunyuanVideo15Transformer3DModel(BaseDiT): A Transformer model for video-like data used in [HunyuanVideo1.5](https://huggingface.co/tencent/HunyuanVideo1.5). """ - # shard single stream, double stream blocks, and refiner_blocks _fsdp_shard_conditions = HunyuanVideo15Config()._fsdp_shard_conditions _compile_conditions = HunyuanVideo15Config()._compile_conditions _supported_attention_backends = HunyuanVideo15Config()._supported_attention_backends @@ -517,13 +504,6 @@ class SingleTokenRefiner(nn.Module): ) -> None: super().__init__() - # Input projection - # self.input_embedder = ReplicatedLinear( - # in_channels, - # hidden_size, - # bias=True, - # params_dtype=dtype, - # prefix=f"{prefix}.input_embedder") self.input_embedder = nn.Linear(in_channels, hidden_size, bias=True) # Timestep embedding @@ -549,7 +529,6 @@ class SingleTokenRefiner(nn.Module): ]) def forward(self, x, t, mask=None): - # Get timestep embeddings timestep_aware_representations = self.t_embedder(t) original_dtype = x.dtype @@ -563,7 +542,6 @@ class SingleTokenRefiner(nn.Module): c = timestep_aware_representations + context_aware_representations # Project input x = self.input_embedder(x) - # Process through refiner blocks for block in self.refiner_blocks: x = block(x, c, mask=mask) return x @@ -617,13 +595,11 @@ class IndividualTokenRefinerBlock(nn.Module): self.attn = LocalAttention( num_heads=num_attention_heads, head_size=hidden_size // num_attention_heads, - # TODO: remove hardcode supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), ) def forward(self, x, c, mask=None): - # Get modulation parameters gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=-1) # Self-attention norm_x = self.norm1(x) @@ -640,7 +616,6 @@ class IndividualTokenRefinerBlock(nn.Module): current_timestep=0, attn_mask=mask, ) - # Run distributed attention with set_forward_context(current_timestep=0, attn_metadata=attn_metadata): attn_output = self.attn(q, k, v) # [B, L, H, D] attn_output = attn_output.reshape(batch_size, seq_len, -1) diff --git a/fastvideo/models/dits/hyworld/__init__.py b/fastvideo/models/dits/hyworld/__init__.py index 5213bf913..b7095f699 100644 --- a/fastvideo/models/dits/hyworld/__init__.py +++ b/fastvideo/models/dits/hyworld/__init__.py @@ -15,7 +15,6 @@ from .resolution_utils import ( get_resolution_from_image, ) __all__ = [ - # Model (used by model registry) "HYWorldTransformer3DModel", "HYWorldDoubleStreamBlock", # Inference utilities (used by examples) diff --git a/fastvideo/models/dits/hyworld/camera_rope.py b/fastvideo/models/dits/hyworld/camera_rope.py index 8f1463e9c..ae9d3077d 100644 --- a/fastvideo/models/dits/hyworld/camera_rope.py +++ b/fastvideo/models/dits/hyworld/camera_rope.py @@ -94,7 +94,6 @@ def prope_qkv( assert q.shape == k.shape == v.shape assert viewmats.shape == (batch, cameras, 4, 4) assert Ks is None or Ks.shape == (batch, cameras, 3, 3) - # assert seqlen == cameras * patches_x * patches_y apply_fn_q, apply_fn_kv, apply_fn_o = _prepare_apply_fns_all_dim( head_dim=head_dim, @@ -134,7 +133,6 @@ def _prepare_apply_fns_all_dim( device = viewmats.device (batch, cameras, _, _) = viewmats.shape - # Normalize camera intrinsics. if Ks is not None: Ks_norm = torch.zeros_like(Ks) Ks_norm[..., 0, 0] = Ks[..., 0, 0] diff --git a/fastvideo/models/dits/hyworld/hyworld.py b/fastvideo/models/dits/hyworld/hyworld.py index 91534d1af..26772734c 100644 --- a/fastvideo/models/dits/hyworld/hyworld.py +++ b/fastvideo/models/dits/hyworld/hyworld.py @@ -65,7 +65,6 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): ) self.hidden_size = hidden_size - # Add ProPE projection layer for camera-aware attention self.img_attn_prope_proj = ReplicatedLinear(hidden_size, hidden_size, bias=True, @@ -104,7 +103,6 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): Returns: Tuple of (img, txt) output tokens """ - # Process modulation vectors (inherited from parent) img_mod_outputs = self.img_mod(vec) ( img_attn_shift, @@ -125,9 +123,7 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): txt_mlp_gate, ) = torch.chunk(txt_mod_outputs, 6, dim=-1) - # Prepare image for attention using fused operation img_attn_input = self.img_attn_norm(img, img_attn_shift, img_attn_scale, convert_modulation_dtype=True) - # Get QKV for image img_qkv, _ = self.img_attn_qkv(img_attn_input) batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1] @@ -135,21 +131,17 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): img_qkv = img_qkv.view(batch_size, image_seq_len, 3, self.num_attention_heads, -1) img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, 2] - # Apply QK-Norm if needed img_q = self.img_attn_q_norm(img_q).to(img_v) img_k = self.img_attn_k_norm(img_k).to(img_v) - # Prepare text for attention using fused operation txt_attn_input = self.txt_attn_norm(txt, txt_attn_shift, txt_attn_scale, convert_modulation_dtype=True) - # Get QKV for text txt_qkv, _ = self.txt_attn_qkv(txt_attn_input) batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1] # Split QKV txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3, self.num_attention_heads, -1) txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, 2] - # Apply QK-Norm if needed txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype) txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype) @@ -174,7 +166,6 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): current_timestep=0, attn_mask=encoder_attention_mask, ) - # Run distributed attention with set_forward_context(current_timestep=0, attn_metadata=attn_metadata): img_attn, txt_attn = self.attn( img_q, @@ -187,7 +178,6 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): freqs_cis=freqs_cis, ) - # begin hyworld # attention with prope (same text mask, so reuse the metadata) # NOTE: Do NOT pass freqs_cis to prope attention - HY-WorldPlay does not apply RoPE to prope with set_forward_context(current_timestep=0, attn_metadata=attn_metadata): @@ -206,13 +196,11 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): img_attn_prope = apply_fn_o(img_attn_prope) # [batch, num_heads, seqlen, head_dim] img_attn_prope = rearrange(img_attn_prope, "B H L D -> B L (H D)") - # add prope to img_attn img_attn_out, _ = self.img_attn_proj(img_attn.view(batch_size, image_seq_len, -1)) img_attn_prope_out, _ = self.img_attn_prope_proj(img_attn_prope) img_attn_out = img_attn_out + img_attn_prope_out # end hyworld - # Use fused operation for residual connection, normalization, and modulation img_mlp_input, img_residual = self.img_attn_residual_mlp_norm(img, img_attn_out, img_attn_gate, @@ -220,14 +208,11 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): img_mlp_scale, convert_modulation_dtype=True) - # Process image MLP img_mlp_out = self.img_mlp(img_mlp_input) img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate) - # Process text attention output txt_attn_out, _ = self.txt_attn_proj(txt_attn.reshape(batch_size, text_seq_len, -1)) - # Use fused operation for residual connection, normalization, and modulation txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm(txt, txt_attn_out, txt_attn_gate, @@ -235,7 +220,6 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): txt_mlp_scale, convert_modulation_dtype=True) - # Process text MLP txt_mlp_out = self.txt_mlp(txt_mlp_input) txt = self.txt_mlp_residual(txt_residual, txt_mlp_out, txt_mlp_gate) @@ -314,7 +298,6 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): for i in range(config.arch_config.num_layers) ]) - # Add action conditioning module self.action_in = TimestepEmbedder(self.hidden_size, act_layer="silu", dtype=None, @@ -382,7 +365,6 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): temb = self.time_in(timestep, timestep_r=timestep_r) temb_txt = self.time_in(timestep_txt, timestep_r=timestep_r) - # Add action conditioning if provided # temb shape: [B*T, C] where T = num_frames temb = temb + self.action_in(action.reshape(-1)) @@ -472,7 +454,6 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): ], dim=0, )) - # Apply same reordering to attention masks new_encoder_attention_mask.append( torch.cat( [ @@ -522,7 +503,6 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): # Apply final_layer on sharded data first, then gather hidden_states = self.final_layer(hidden_states, temb) - # Gather the output from all ranks if get_sp_world_size() > 1: hidden_states = hidden_states.contiguous() hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1) diff --git a/fastvideo/models/dits/hyworld/pose.py b/fastvideo/models/dits/hyworld/pose.py index 40410591f..5edb7b3b8 100644 --- a/fastvideo/models/dits/hyworld/pose.py +++ b/fastvideo/models/dits/hyworld/pose.py @@ -95,7 +95,6 @@ def parse_pose_string( num_frames = int(duration) - # Parse action and create motion dicts if action == "w": # Forward for _ in range(num_frames): @@ -184,14 +183,11 @@ def pose_to_input( - intrinsics: Normalized camera intrinsics [T, 3, 3] - action_labels: Action labels for each frame [T] """ - # Handle different input types if isinstance(pose_data, str): if pose_data.endswith(".json"): - # Load from JSON file with open(pose_data, "r") as f: pose_json = json.load(f) else: - # Parse pose string pose_json = pose_string_to_json(pose_data) elif isinstance(pose_data, dict): pose_json = pose_data @@ -212,7 +208,6 @@ def pose_to_input( w2c = np.linalg.inv(c2w) w2c_list.append(w2c) - # Normalize intrinsics intrinsic = np.array(pose_json[t_key]["K"]) intrinsic[0, 0] /= intrinsic[0, 2] * 2 intrinsic[1, 1] /= intrinsic[1, 2] * 2 @@ -223,14 +218,12 @@ def pose_to_input( w2c_list = np.array(w2c_list) intrinsic_list = torch.tensor(np.array(intrinsic_list)) - # Compute relative camera-to-world transforms c2ws = np.linalg.inv(w2c_list) C_inv = np.linalg.inv(c2ws[:-1]) relative_c2w = np.zeros_like(c2ws) relative_c2w[0, ...] = c2ws[0, ...] relative_c2w[1:, ...] = C_inv @ c2ws[1:, ...] - # Initialize one-hot action encodings trans_one_hot = np.zeros((relative_c2w.shape[0], 4), dtype=np.int32) rotate_one_hot = np.zeros((relative_c2w.shape[0], 4), dtype=np.int32) @@ -276,7 +269,6 @@ def pose_to_input( trans_one_hot = torch.tensor(trans_one_hot) rotate_one_hot = torch.tensor(rotate_one_hot) - # Convert one-hot to single-dimension labels trans_one_label = one_hot_to_one_dimension(trans_one_hot) rotate_one_label = one_hot_to_one_dimension(rotate_one_hot) action_one_label = trans_one_label * 9 + rotate_one_label @@ -363,7 +355,6 @@ def parse_pose_string_to_actions(pose_string: str, fps: int = 24) -> list[dict]: else: raise ValueError(f"Unknown action: {action}") - # Add frame-level actions for _ in range(num_frames): frame_actions.append(action_values.copy()) diff --git a/fastvideo/models/dits/hyworld/resolution_utils.py b/fastvideo/models/dits/hyworld/resolution_utils.py index 737c45f81..9d1382207 100644 --- a/fastvideo/models/dits/hyworld/resolution_utils.py +++ b/fastvideo/models/dits/hyworld/resolution_utils.py @@ -64,7 +64,6 @@ def get_resolution_from_image(image_path, target_resolution="480p"): Returns: tuple[int, int]: (height, width) matching HY-WorldPlay's bucket selection """ - # Handle URL inputs if isinstance(image_path, str) and image_path.startswith(('http://', 'https://')): response = requests.get(image_path) response.raise_for_status() diff --git a/fastvideo/models/dits/hyworld/retrieval_context.py b/fastvideo/models/dits/hyworld/retrieval_context.py index e5ccf3343..8947e4b5c 100644 --- a/fastvideo/models/dits/hyworld/retrieval_context.py +++ b/fastvideo/models/dits/hyworld/retrieval_context.py @@ -76,7 +76,6 @@ def rotation_matrix_to_angles(R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tens y = fwd[1] z = fwd[2] - # compute yaw and pitch yaw_rad = torch.atan2(x, z) yaw_deg = yaw_rad * (180.0 / math.pi) pitch_rad = torch.atan2(y, torch.sqrt(x**2 + z**2)) @@ -121,17 +120,14 @@ def is_inside_fov_3d_hv( # Calculate the horizontal angle (yaw/azimuth), assuming the Z-axis is forward. azimuth = torch.atan2(x, z) * (180 / math.pi) - # Calculate the vertical angle (pitch/elevation). elevation = torch.atan2(y, torch.sqrt(x**2 + z**2)) * (180 / math.pi) - # Calculate the angular difference from the center view direction (handling angle wrapping). diff_azimuth = azimuth - center_yaw diff_azimuth = torch.remainder(diff_azimuth + 180, 360) - 180 diff_elevation = elevation - center_pitch diff_elevation = torch.remainder(diff_elevation + 180, 360) - 180 - # Check if within FOV in_fov_h = diff_azimuth.abs() < fov_half_h in_fov_v = diff_elevation.abs() < fov_half_v @@ -178,7 +174,6 @@ def calculate_fov_overlap_similarity( fov_half_h = torch.tensor(fov_h_deg / 2.0, device=device) fov_half_v = torch.tensor(fov_v_deg / 2.0, device=device) - # move to P_w_curr (N, 3) points_world = points_local + P_w_curr[None, :] in_fov_curr = is_inside_fov_3d_hv( @@ -190,7 +185,6 @@ def calculate_fov_overlap_similarity( fov_half_v, ) - # compute based on angle in_fov_hist = is_inside_fov_3d_hv( points_world, P_w_hist[None, :], @@ -200,7 +194,6 @@ def calculate_fov_overlap_similarity( fov_half_v, ) - # compute based on distance dist = torch.norm(points_world - P_w_hist.reshape(1, -1), dim=1) < 8.0 in_fov_hist = in_fov_hist.bool() & dist.reshape(1, -1).bool() @@ -292,7 +285,6 @@ def select_aligned_memory_frames( candidate_distances.sort(key=lambda x: x[1]) for start_idx, _ in candidate_distances: - # check the memory frame number if len(memory_frames_indices) >= memory_frames: break diff --git a/fastvideo/models/dits/lingbotworld/cam_utils.py b/fastvideo/models/dits/lingbotworld/cam_utils.py index 9190ce27e..5384fbaff 100644 --- a/fastvideo/models/dits/lingbotworld/cam_utils.py +++ b/fastvideo/models/dits/lingbotworld/cam_utils.py @@ -67,7 +67,6 @@ def compute_relative_poses( # ensure identity matrix for 1st frame relative_poses[0] = torch.eye(4, device=c2ws_mat.device, dtype=c2ws_mat.dtype) if framewise: - # compute pose between i and i+1 relative_poses_framewise = torch.bmm(SE3_inverse(relative_poses[:-1]), relative_poses[1:]) relative_poses[1:] = relative_poses_framewise if normalize_trans: # note refer to camctrl2: "we scale the coordinate inputs to roughly 1 standard deviation to simplify model learning." @@ -116,8 +115,6 @@ def get_plucker_embeddings( rays_d = directions @ c2ws_mat[:, :3, :3].transpose(-1, -2) # [f, h*w, 3] rays_o = c2ws_mat[:, :3, 3] # [f, 3] rays_o = rays_o[:, None, :].expand_as(rays_d) # [f, h*w, 3] - # rays_dxo = torch.cross(rays_o, rays_d, dim=-1) # [f, h*w, 3] - # note refer to: apt2 plucker_embeddings = torch.cat([rays_o, rays_d], dim=-1) # [f, h*w, 6] plucker_embeddings = plucker_embeddings.view([n_frames, height, width, 6]) # [f*h*w, 6] return plucker_embeddings diff --git a/fastvideo/models/dits/lingbotworld/model.py b/fastvideo/models/dits/lingbotworld/model.py index e88f7afc7..47de6768a 100644 --- a/fastvideo/models/dits/lingbotworld/model.py +++ b/fastvideo/models/dits/lingbotworld/model.py @@ -13,8 +13,6 @@ from fastvideo.distributed.communication_op import (sequence_model_parallel_all_ from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidual, ScaleResidualLayerNormScaleShift) from fastvideo.layers.linear import ReplicatedLinear -# from torch.nn import RMSNorm -# TODO: RMSNorm .... from fastvideo.layers.mlp import MLP from fastvideo.layers.rotary_embedding import get_rotary_pos_embed from fastvideo.layers.visual_embedding import (PatchEmbed, WanCamControlPatchEmbedding) @@ -138,7 +136,6 @@ class LingBotWorldTransformerBlock(nn.Module): hidden_states = hidden_states.squeeze(1) bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype - # assert orig_dtype != torch.float32 if temb.dim() == 4: # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) @@ -297,7 +294,6 @@ class LingBotWorldTransformer3DModel(BaseDiT): post_patch_height = height // p_h post_patch_width = width // p_w - # Get rotary embeddings d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] freqs_cos, freqs_sin = get_rotary_pos_embed((post_patch_num_frames, post_patch_height, post_patch_width), @@ -320,7 +316,6 @@ class LingBotWorldTransformer3DModel(BaseDiT): # Shard with padding support - returns (sharded_tensor, original_seq_len) hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1) - # Shard c2ws_plucker_emb if c2ws_plucker_emb is not None: c2ws_plucker_emb, _ = sequence_model_parallel_shard(c2ws_plucker_emb, dim=1) @@ -375,7 +370,6 @@ class LingBotWorldTransformer3DModel(BaseDiT): hidden_states = self.norm_out(hidden_states, shift, scale) - # Gather and unpad in one operation hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1) hidden_states = self.proj_out(hidden_states) diff --git a/fastvideo/models/dits/lingbotworld2/cam_utils.py b/fastvideo/models/dits/lingbotworld2/cam_utils.py index 50cb3d018..b0e9c030f 100644 --- a/fastvideo/models/dits/lingbotworld2/cam_utils.py +++ b/fastvideo/models/dits/lingbotworld2/cam_utils.py @@ -65,7 +65,6 @@ def compute_relative_poses( # ensure identity matrix for 1st frame relative_poses[0] = torch.eye(4, device=c2ws_mat.device, dtype=c2ws_mat.dtype) if framewise: - # compute pose between i and i+1 relative_poses_framewise = torch.bmm(SE3_inverse(relative_poses[:-1]), relative_poses[1:]) relative_poses[1:] = relative_poses_framewise if normalize_trans: # note refer to camctrl2: "we scale the coordinate inputs to roughly 1 standard deviation to simplify model learning." diff --git a/fastvideo/models/dits/ltx2.py b/fastvideo/models/dits/ltx2.py index 9bb18a07a..ca0f7bc26 100644 --- a/fastvideo/models/dits/ltx2.py +++ b/fastvideo/models/dits/ltx2.py @@ -1324,7 +1324,6 @@ class LTXDistributedAttention(DistributedAttention): # Transpose back to [2*B, T, H/SP, D] qkv[:batch_size * 2] = qk_part.transpose(1, 2) - # Apply backend-specific preprocess_qkv qkv = self.attn_impl.preprocess_qkv(qkv, ctx_attn_metadata) # Concatenate with replicated QKV if provided @@ -1348,7 +1347,6 @@ class LTXDistributedAttention(DistributedAttention): output = output[:, :split_idx] replicated_output = sequence_model_parallel_all_gather(replicated_output.contiguous(), dim=2) - # Apply backend-specific postprocess_output output = self.attn_impl.postprocess_output(output, ctx_attn_metadata) output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len)) @@ -1409,7 +1407,6 @@ class LTXLocalAttention(LocalAttention): forward_context: ForwardContext = get_forward_context() ctx_attn_metadata = forward_context.attn_metadata - # Apply LTX-2 style RoPE if ltx_freqs_cis is not None: # Use separate K RoPE if provided (for cross-attention), otherwise use same as Q k_freqs = ltx_k_freqs_cis if ltx_k_freqs_cis is not None else ltx_freqs_cis @@ -2167,7 +2164,6 @@ class BasicAVTransformerBlock(torch.nn.Module): # Audio-to-video cross-attention: Q from video (sharded), K/V from audio # For SP, gather audio context from all ranks before cross-attention if self.use_distributed_attention: - # Gather audio context from all SP ranks ax_context = sequence_model_parallel_all_gather(ax_scaled, dim=1) # Video Q: video tokens divide evenly (adjusted upstream), so simple slicing works video_q_pe = None @@ -2207,10 +2203,8 @@ class BasicAVTransformerBlock(torch.nn.Module): # For SP, gather both modalities and run cross-attention on full sequences # (audio is short, so no need to keep it sharded for cross-attention) if self.use_distributed_attention: - # Gather both modalities to full sequences ax_full = sequence_model_parallel_all_gather(ax_scaled, dim=1) vx_context = sequence_model_parallel_all_gather(vx_scaled, dim=1) - # Use full positional embeddings directly audio_q_pe = audio.cross_positional_embeddings video_k_pe = video.cross_positional_embeddings # Trim gathered tensors to original lengths (may have padding from sharding) @@ -2220,7 +2214,6 @@ class BasicAVTransformerBlock(torch.nn.Module): if video_k_pe is not None: original_video_len = video_k_pe[0].shape[2] vx_context = vx_context[:, :original_video_len, :] - # Run cross-attention on full audio v2a_out_full = self.video_to_audio_attn( ax_full, context=vx_context, @@ -2732,12 +2725,10 @@ class LTX2Transformer3DModel(BaseDiT): arch = config.arch_config - # Get SP world size for distributed attention sp_world_size = get_sp_world_size() use_vsa_backend = effective_attention_backend(config) == AttentionBackendEnum.VIDEO_SPARSE_ATTN use_distributed_attention = sp_world_size > 1 or use_vsa_backend - # Validate that attention heads are divisible by SP world size if sp_world_size > 1: assert arch.num_attention_heads % sp_world_size == 0, ( f"The number of video attention heads ({arch.num_attention_heads}) " @@ -2824,10 +2815,8 @@ class LTX2Transformer3DModel(BaseDiT): ) -> torch.Tensor: if isinstance(encoder_hidden_states, list): encoder_hidden_states = encoder_hidden_states[0] - # Get SP parameters sp_world_size = get_sp_world_size() - # Get fps for position computation fps = None try: forward_ctx = get_forward_context() @@ -2850,7 +2839,6 @@ class LTX2Transformer3DModel(BaseDiT): if video_sigma is None: video_sigma = timestep[:, 0] if timestep.ndim > 1 else timestep - # Shard video latents and timestep across SP ranks video_original_seq_len = latents.shape[1] video_padded_seq_len = video_original_seq_len video_timestep = timestep @@ -2905,7 +2893,6 @@ class LTX2Transformer3DModel(BaseDiT): keyframes_mask=keyframes_mask, ) - # Process audio modality if provided audio_modality = None audio_shape = None audio_original_seq_len = 0 @@ -2916,12 +2903,10 @@ class LTX2Transformer3DModel(BaseDiT): if audio_sigma is None: audio_sigma = (audio_timestep[:, 0] if audio_timestep.ndim > 1 else audio_timestep) - # Shard audio latents and timestep across SP ranks audio_original_seq_len = audio_latents.shape[1] sharded_audio_timestep = audio_timestep if sp_world_size > 1: audio_latents, audio_original_seq_len = sequence_model_parallel_shard(audio_latents, dim=1) - # Shard audio timestep along sequence dimension sharded_audio_timestep, _ = sequence_model_parallel_shard(audio_timestep, dim=1) # Compute audio RoPE positions for the FULL sequence (before sharding) @@ -2939,7 +2924,6 @@ class LTX2Transformer3DModel(BaseDiT): sigma=audio_sigma, ) - # Run transformer with original sequence lengths video_out, audio_out = self.model( video=video_modality, audio=audio_modality, @@ -2964,11 +2948,9 @@ class LTX2Transformer3DModel(BaseDiT): audio_modality.timesteps, ) - # Gather and unpad video output if sp_world_size > 1 and video_out is not None: video_out = sequence_model_parallel_all_gather_with_unpad(video_out, video_original_seq_len, dim=1) - # Gather and unpad audio output if sp_world_size > 1 and audio_out is not None: audio_out = sequence_model_parallel_all_gather_with_unpad(audio_out, audio_original_seq_len, dim=1) diff --git a/fastvideo/models/dits/magi_human.py b/fastvideo/models/dits/magi_human.py index 88d7d13bf..3bbc4867a 100644 --- a/fastvideo/models/dits/magi_human.py +++ b/fastvideo/models/dits/magi_human.py @@ -59,9 +59,6 @@ from fastvideo.layers.rotary_embedding import _apply_rotary_emb from fastvideo.models.dits.base import BaseDiT from fastvideo.platforms import AttentionBackendEnum -# --------------------------------------------------------------------------- -# Enums -# --------------------------------------------------------------------------- class Modality(IntEnum): @@ -70,9 +67,6 @@ class Modality(IntEnum): TEXT = 2 -# --------------------------------------------------------------------------- -# Activations -# --------------------------------------------------------------------------- def swiglu7(x: torch.Tensor, alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor: @@ -92,9 +86,6 @@ def gelu7(x: torch.Tensor, alpha: float = 1.702, limit: float = 7.0) -> torch.Te return (x * torch.sigmoid(alpha * x)).to(in_dtype) -# --------------------------------------------------------------------------- -# Modality dispatcher -# --------------------------------------------------------------------------- class ModalityDispatcher: @@ -130,9 +121,6 @@ class ModalityDispatcher: return x[inv_permute_mapping] -# --------------------------------------------------------------------------- -# Norms, rotary embed -# --------------------------------------------------------------------------- class MultiModalityRMSNorm(nn.Module): @@ -218,9 +206,6 @@ class ElementWiseFourierEmbed(nn.Module): return torch.cat((sin_proj, cos_proj), dim=1).flatten(1) -# --------------------------------------------------------------------------- -# Packed-expert linear -# --------------------------------------------------------------------------- class PackedExpertLinear(nn.Module): @@ -277,9 +262,6 @@ class PackedExpertLinear(nn.Module): return modality_dispatcher.undispatch(*parts) -# --------------------------------------------------------------------------- -# Attention & MLP -# --------------------------------------------------------------------------- @dataclass @@ -620,9 +602,6 @@ class MagiTransformerLayer(nn.Module): return hidden_states + mlp_out -# --------------------------------------------------------------------------- -# Adapter (per-modality embedders + Fourier RoPE producer) -# --------------------------------------------------------------------------- class MagiAdapter(nn.Module): @@ -691,9 +670,6 @@ class MagiAdapter(nn.Module): return out, rope -# --------------------------------------------------------------------------- -# Top-level DiT -# --------------------------------------------------------------------------- class _TransformerBlock(nn.Module): diff --git a/fastvideo/models/dits/matrixgame2/action_module.py b/fastvideo/models/dits/matrixgame2/action_module.py index 4b2558178..c74a2039f 100644 --- a/fastvideo/models/dits/matrixgame2/action_module.py +++ b/fastvideo/models/dits/matrixgame2/action_module.py @@ -35,7 +35,6 @@ def _get_nd_rotary_pos_embed_matrixgame2( theta_rescale_factor=theta_rescale_factor, dtype=torch.float32, ) - # convert from [S, D/2] to [S, D] format cos = cos.repeat_interleave(2, dim=1) sin = sin.repeat_interleave(2, dim=1) return cos, sin @@ -54,11 +53,9 @@ def _apply_rotary_emb_qk( cos = freqs_cos[start_offset:start_offset + seq_len] # [S, D] sin = freqs_sin[start_offset:start_offset + seq_len] # [S, D] - # Move to device cos = cos.to(xq.device) sin = sin.to(xq.device) - # Convert from [S, D] (interleaved) back to [S, D/2] cos_half = cos[:, ::2] # [S, D/2] sin_half = sin[:, ::2] # [S, D/2] @@ -160,15 +157,12 @@ def _update_kv_cache_and_attend( :, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens, ].clone()) - # Calculate indices with eviction adjustment local_end_index = (original_local_end_index + current_end - original_global_end_index - num_evicted_tokens) local_start_index = local_end_index - num_new_tokens else: - # Calculate indices without eviction local_end_index = (original_local_end_index + current_end - original_global_end_index) local_start_index = local_end_index - num_new_tokens - # Store new k, v in cache if store_first_only: kv_cache["k"][:, local_start_index:local_end_index] = k[:kv_cache["k"].shape[0]] kv_cache["v"][:, local_start_index:local_end_index] = v[:kv_cache["v"].shape[0]] @@ -189,7 +183,6 @@ def _update_kv_cache_and_attend( attn = attn_layer(q, cached_k, cached_v) - # Update indices if isinstance(kv_cache["global_end_index"], torch.Tensor): kv_cache["global_end_index"].fill_(current_end) else: @@ -231,7 +224,6 @@ class ActionModule(nn.Module): blocks: list | None = None, ): super().__init__() - # Initialize mutable defaults patch_size = patch_size if patch_size is not None else [1, 2, 2] rope_dim_list = (rope_dim_list if rope_dim_list is not None else [8, 28, 28]) mouse_qk_dim_list = (mouse_qk_dim_list if mouse_qk_dim_list is not None else [8, 28, 28]) @@ -369,7 +361,6 @@ class ActionModule(nn.Module): if rope_dim_list is None: rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)] assert sum(rope_dim_list) == head_dim, ("sum(rope_dim_list) should equal to head_dim of attention layer") - # Use Matrix-Game wrapper for FastVideo's function freqs_cos, freqs_sin = _get_nd_rotary_pos_embed_matrixgame2( rope_dim_list, rope_sizes, @@ -442,7 +433,6 @@ class ActionModule(nn.Module): k = self.img_attn_k_norm(k).to(v) # rope embd - # freqs_cis = (self.freqs_cos, self.freqs_sin) q, k = _apply_rotary_emb_qk(q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame) ## TODO: adding cache here @@ -474,8 +464,6 @@ class ActionModule(nn.Module): ) else: attn = self.mouse_attn_layer(q, k, v) - # Compute cu_squlens and max_seqlen for flash attention - # qk norm attn = rearrange(attn, "(b S) T h d -> b (T S) (h d)", b=B) attn, _ = self.proj_mouse(attn) return attn @@ -529,7 +517,6 @@ class ActionModule(nn.Module): ] group_keyboard = torch.stack(group_keyboard, dim=1) # B F RW C group_keyboard = group_keyboard.reshape(shape=(group_keyboard.shape[0], group_keyboard.shape[1], -1)) - # apply cross attn mouse_q, _ = self.mouse_attn_q(hidden_states) keyboard_kv, _ = self.keyboard_attn_kv(group_keyboard) @@ -540,14 +527,10 @@ class ActionModule(nn.Module): B, L, KHD = keyboard_kv.shape k, v = keyboard_kv.view(B, L, 2, self.heads_num, D).permute(2, 0, 1, 3, 4) - # Compute cu_squlens and max_seqlen for flash attention - # qk norm q = self.key_attn_q_norm(q).to(v) k = self.key_attn_k_norm(k).to(v) S = th * tw - # assert S == 880 - # position embed if use_rope_keyboard: B, TS, H, D = q.shape T_ = TS // S diff --git a/fastvideo/models/dits/matrixgame2/causal_model.py b/fastvideo/models/dits/matrixgame2/causal_model.py index 54d6b0883..b604ee460 100644 --- a/fastvideo/models/dits/matrixgame2/causal_model.py +++ b/fastvideo/models/dits/matrixgame2/causal_model.py @@ -55,7 +55,6 @@ def causal_rope_apply(x: torch.Tensor, grid_sizes: tuple[int, int, int], freqs: # split freqs freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) - # loop over samples output = [] f, h, w = grid_sizes @@ -73,11 +72,9 @@ def causal_rope_apply(x: torch.Tensor, grid_sizes: tuple[int, int, int], freqs: dim=-1, ).reshape(seq_len, 1, -1) - # apply rotary embedding x_i = torch.view_as_real(x_i * freqs_i).flatten(2) x_i = torch.cat([x_i, x[i, seq_len:]]) - # append to collection output.append(x_i) return torch.stack(output).type_as(x) @@ -178,7 +175,6 @@ class CausalMatrixGame2SelfAttention(nn.Module): if cache_start is None: cache_start = current_start - # Calculate start_frame for causal mode if kv_cache is not None: frame_seqlen = int(grid_sizes[1] * grid_sizes[2]) start_frame = current_start // frame_seqlen @@ -249,7 +245,6 @@ class CausalMatrixGame2SelfAttention(nn.Module): block_mask=block_mask, )[:, :, :-padded_length].transpose(2, 1) else: - # Calculate frame_seqlen correctly from grid_sizes (single frame token count) if grid_sizes is not None: frame_seqlen = int(grid_sizes[1] * grid_sizes[2]) else: @@ -260,7 +255,6 @@ class CausalMatrixGame2SelfAttention(nn.Module): current_end = current_start + q.shape[1] sink_tokens = self.sink_size * frame_seqlen - # Compute max_attention_size dynamically based on actual frame_seqlen max_attention_size = (15 * frame_seqlen if self.local_attn_size == -1 else self.local_attn_size * frame_seqlen) @@ -465,7 +459,6 @@ class CausalMatrixGame2TransformerBlock(nn.Module): frame_seqlen = hidden_states.shape[1] // num_frames bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype - # hidden_states = hidden_states.float() e = self.scale_shift_table + temb assert e.shape == (bs, num_frames, 6, self.hidden_dim) @@ -869,7 +862,6 @@ class CausalMatrixGame2WanModel(BaseDiT): start_frame: int = 0, **kwargs, ) -> torch.Tensor: - # Extract num_frame_per_block from kwargs effective_num_frame_per_block = kwargs.pop("num_frame_per_block", self.num_frame_per_block) if mouse_cond is not None or keyboard_cond is not None: @@ -1212,7 +1204,6 @@ class CausalMatrixGame2WanModel(BaseDiT): num_frame_per_block=self.num_frame_per_block, ) - # Initialize kv_cache, kv_cache_mouse, kv_cache_keyboard, crossattn_cache if not provided kv_cache = kwargs.pop("kv_cache", [None] * len(self.blocks)) kv_cache_mouse = kwargs.pop("kv_cache_mouse", [None] * len(self.blocks)) kv_cache_keyboard = kwargs.pop("kv_cache_keyboard", [None] * len(self.blocks)) diff --git a/fastvideo/models/dits/matrixgame2/model.py b/fastvideo/models/dits/matrixgame2/model.py index e53bbbb73..def119714 100644 --- a/fastvideo/models/dits/matrixgame2/model.py +++ b/fastvideo/models/dits/matrixgame2/model.py @@ -37,7 +37,6 @@ from fastvideo.models.wan.transformer import ( ) from fastvideo.platforms import AttentionBackendEnum, current_platform -# Import ActionModule from .action_module import ActionModule logger = init_logger(__name__) @@ -102,7 +101,6 @@ class MatrixGame2CrossAttention(WanSelfAttention): """ b, n, d = x.size(0), self.num_heads, self.head_dim - # compute query, key, value q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) if crossattn_cache is not None: @@ -119,7 +117,6 @@ class MatrixGame2CrossAttention(WanSelfAttention): k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) v = self.to_v(context)[0].view(b, -1, n, d) - # compute attention x = self.attn(q, k, v) # output @@ -273,7 +270,6 @@ class MatrixGame2TransformerBlock(nn.Module): key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - # Apply rotary embeddings cos, sin = freqs_cis query, key = ( _apply_rotary_emb(query, cos, sin, is_neox_style=False), @@ -316,7 +312,6 @@ class MatrixGame2TransformerBlock(nn.Module): keyboard_cond, num_frame_per_block=int(grid_sizes[0]), ) - # ================================================= # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) @@ -428,7 +423,6 @@ class MatrixGame2WanModel(BaseDiT): post_patch_height = height // p_h post_patch_width = width // p_w - # Get rotary embeddings d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] freqs_cos, freqs_sin = get_rotary_pos_embed( diff --git a/fastvideo/models/dits/matrixgame2/utils.py b/fastvideo/models/dits/matrixgame2/utils.py index 71f6981a2..1304f9173 100644 --- a/fastvideo/models/dits/matrixgame2/utils.py +++ b/fastvideo/models/dits/matrixgame2/utils.py @@ -408,13 +408,11 @@ def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = No if num_frames % 4 != 1: raise ValueError("Matrix-Game conditioning expects num_frames to be 4k+1.") - # Set seed for reproducibility if provided if seed is not None: random.seed(seed) num_samples_per_action = 4 - # Define actions based on keyboard_dim if keyboard_dim == 4: # Universal model: W, S, A, D actions_single_action = ["forward", "left", "right"] @@ -444,7 +442,6 @@ def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = No for camera in actions_single_camera: actions_to_test.append(f"{action}_{camera}") - # Ensure we have at least some actions if not actions_to_test: actions_to_test = actions_single_action * 5 diff --git a/fastvideo/models/dits/matrixgame3/action_module.py b/fastvideo/models/dits/matrixgame3/action_module.py index bd1fa4aec..40235e925 100644 --- a/fastvideo/models/dits/matrixgame3/action_module.py +++ b/fastvideo/models/dits/matrixgame3/action_module.py @@ -40,7 +40,6 @@ def _get_nd_rotary_pos_embed_matrixgame( theta_rescale_factor=theta_rescale_factor, dtype=torch.float32, ) - # convert from [S, D/2] to [S, D] format cos = cos.repeat_interleave(2, dim=1) sin = sin.repeat_interleave(2, dim=1) return cos, sin @@ -223,7 +222,6 @@ class MatrixGame3ActionModule(nn.Module): if rope_dim_list is None: rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)] assert sum(rope_dim_list) == head_dim, ("sum(rope_dim_list) should equal to head_dim of attention layer") - # Use Matrix-Game wrapper for FastVideo's function freqs_cos, freqs_sin = _get_nd_rotary_pos_embed_matrixgame( rope_dim_list, rope_sizes, diff --git a/fastvideo/models/dits/matrixgame3/model.py b/fastvideo/models/dits/matrixgame3/model.py index fa29fac5f..2408d8dad 100644 --- a/fastvideo/models/dits/matrixgame3/model.py +++ b/fastvideo/models/dits/matrixgame3/model.py @@ -28,7 +28,6 @@ from fastvideo.models.dits.base import BaseDiT from fastvideo.models.wan.transformer import WanSelfAttention from fastvideo.platforms import AttentionBackendEnum -# Import ActionModule from .action_module import MatrixGame3ActionModule logger = init_logger(__name__) @@ -348,7 +347,6 @@ class MatrixGame3TransformerBlock(nn.Module): key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - # Apply rotary embeddings grid_frames, grid_height, grid_width = grid_sizes if memory_length > 0: @@ -447,7 +445,6 @@ class MatrixGame3TransformerBlock(nn.Module): norm_hidden_states = self.cross_attn_residual_norm.norm(hidden_states.float()) norm_hidden_states = norm_hidden_states * (1 + c_scale_msa) + c_shift_msa - # ================================================= # 3. Feed-forward ff_output = self.ffn(norm_hidden_states.to(self.ffn.fc_in.weight.dtype)) diff --git a/fastvideo/models/encoders/clip.py b/fastvideo/models/encoders/clip.py index 9bf28014e..a9e69b115 100644 --- a/fastvideo/models/encoders/clip.py +++ b/fastvideo/models/encoders/clip.py @@ -8,7 +8,6 @@ from collections.abc import Iterable import torch import torch.nn as nn -# from transformers.modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask from fastvideo.attention import LocalAttention from fastvideo.configs.models.encoders import (BaseEncoderOutput, CLIPTextConfig, CLIPVisionConfig) from fastvideo.distributed import divide, get_tp_world_size @@ -18,8 +17,6 @@ from fastvideo.layers.quantization import QuantizationConfig from fastvideo.logger import init_logger from fastvideo.models.encoders.base import ImageEncoder, TextEncoder from fastvideo.models.encoders.vision import resolve_visual_encoder_outputs -# TODO: support quantization -# from vllm.model_executor.layers.quantization import QuantizationConfig from fastvideo.models.loader.weight_utils import default_weight_loader logger = init_logger(__name__) @@ -167,7 +164,6 @@ class CLIPAttention(nn.Module): qkv_states, _ = self.qkv_proj(hidden_states) query_states, key_states, value_states = qkv_states.chunk(3, dim=-1) - # use flash_attn_func query_states = query_states.reshape(query_states.shape[0], query_states.shape[1], self.num_heads_per_partition, self.head_dim) key_states = key_states.reshape(key_states.shape[0], key_states.shape[1], self.num_heads_per_partition, @@ -342,19 +338,10 @@ class CLIPTextTransformer(nn.Module): # input_shape, hidden_states.dtype, device=hidden_states.device # ) - # # expand attention_mask - # if attention_mask is not None and not self._use_flash_attention_2: - # raise NotImplementedError("attention_mask is not supported for CLIPTextTransformer") - # # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] - # attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) encoder_outputs = self.encoder( inputs_embeds=hidden_states, - # attention_mask=attention_mask, - # causal_attention_mask=causal_attention_mask, - # output_attentions=output_attentions, return_all_hidden_states=output_hidden_states, - # return_dict=return_dict, ) last_hidden_state = encoder_outputs[-1] @@ -363,7 +350,6 @@ class CLIPTextTransformer(nn.Module): if self.eos_token_id == 2: # The `eos_token_id` was incorrect before PR #24773: Let's keep what have been done here. # A CLIP model with such `eos_token_id` in the config can't work correctly with extra new tokens added - # ------------------------------------------------------------ # text_embeds.shape = [batch_size, sequence_length, transformer.width] # take features from the eot embedding (eot_token is the highest number in each sequence) # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14 @@ -385,7 +371,6 @@ class CLIPTextTransformer(nn.Module): last_hidden_state=last_hidden_state, pooler_output=pooled_output, hidden_states=encoder_outputs, - # attentions=encoder_outputs.attentions, ) @@ -418,7 +403,6 @@ class CLIPTextModel(TextEncoder): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Define mapping for stacked parameters stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), @@ -440,7 +424,6 @@ class CLIPTextModel(TextEncoder): for name, loaded_weight in weights: if needs_text_model_prefix: name = f"text_model.{name}" - # Handle q_proj, k_proj, v_proj -> qkv_proj mapping for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name in name: # Replace the weight name with the parameter name @@ -453,7 +436,6 @@ class CLIPTextModel(TextEncoder): loaded_params.add(model_param_name) break else: - # Use default weight loader for all other parameters if name in params_dict: param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) @@ -558,7 +540,6 @@ class CLIPVisionTransformer(nn.Module): if not return_all_hidden_states: encoder_outputs = encoder_outputs[0] - # Handle post-norm (if applicable) and stacks feature layers if needed encoder_outputs = resolve_visual_encoder_outputs(encoder_outputs, feature_sample_layers, self.post_layernorm, self.config.num_hidden_layers) diff --git a/fastvideo/models/encoders/llama.py b/fastvideo/models/encoders/llama.py index f04867bc5..90dc44d22 100644 --- a/fastvideo/models/encoders/llama.py +++ b/fastvideo/models/encoders/llama.py @@ -29,9 +29,7 @@ from typing import Any import torch from torch import nn -# from vllm.model_executor.layers.quantization import QuantizationConfig from fastvideo.attention import LocalAttention -# from ..utils import (extract_layer_index) from fastvideo.configs.models.encoders import BaseEncoderOutput, LlamaConfig from fastvideo.distributed import get_tp_world_size from fastvideo.layers.activation import SiluAndMul @@ -59,7 +57,6 @@ class LlamaMLP(nn.Module): self.gate_up_proj = MergedColumnParallelLinear( input_size=hidden_size, output_sizes=[intermediate_size] * 2, - # output_size=intermediate_size, bias=bias, quant_config=quant_config, prefix=f"{prefix}.gate_up_proj", @@ -98,7 +95,6 @@ class LlamaAttention(nn.Module): bias_o_proj: bool = False, prefix: str = "") -> None: super().__init__() - # layer_idx = extract_layer_index(prefix) self.hidden_size = hidden_size tp_size = get_tp_world_size() self.total_num_heads = num_heads @@ -172,18 +168,11 @@ class LlamaAttention(nn.Module): qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) q, k = self.rotary_emb(positions, q, k) - # attn_output = self.attn(q, k, v) - # use flash_attn_func - # TODO (Attn abstraction and backend) - # from flash_attn import flash_attn_func - # reshape q, k, v to (batch_size, seq_len, num_heads, head_dim) batch_size = q.shape[0] seq_len = q.shape[1] q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim) k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) - # import pdb; pdb.set_trace() - # attn_output = flash_attn_func(q, k, v, softmax_scale=self.scaling, causal=True) attn_output = self.attn(q, k, v) attn_output = attn_output.reshape(batch_size, seq_len, self.num_heads * self.head_dim) @@ -319,23 +308,17 @@ class LlamaModel(TextEncoder): all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None for layer in self.layers: if all_hidden_states is not None: - # TODO all_hidden_states += (hidden_states, ) if residual is None else (hidden_states + residual, ) hidden_states, residual = layer(position_ids, hidden_states, residual) hidden_states, _ = self.norm(hidden_states, residual) - # add hidden states from the last decoder layer if all_hidden_states is not None: all_hidden_states += (hidden_states, ) - # TODO(will): maybe unify the output format with other models and use - # our own class output = BaseEncoderOutput( last_hidden_state=hidden_states, - # past_key_values=past_key_values if use_cache else None, hidden_states=all_hidden_states, - # attentions=all_self_attns, ) return output @@ -351,17 +334,6 @@ class LlamaModel(TextEncoder): # Models trained using ColossalAI may include these tensors in # the checkpoint. Skip them. continue - # if (self.quant_config is not None and - # (scale_name := self.quant_config.get_cache_scale(name))): - # # Loading kv cache quantization scales - # param = params_dict[scale_name] - # weight_loader = getattr(param, "weight_loader", - # default_weight_loader) - # loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else - # loaded_weight[0]) - # weight_loader(param, loaded_weight) - # loaded_params.add(scale_name) - # continue if "scale" in name: # Remapping the name of FP8 kv-scale. kv_scale_name: str | None = maybe_remap_kv_scale_name(name, params_dict) diff --git a/fastvideo/models/encoders/qwen2_5.py b/fastvideo/models/encoders/qwen2_5.py index a94840c18..cffded836 100644 --- a/fastvideo/models/encoders/qwen2_5.py +++ b/fastvideo/models/encoders/qwen2_5.py @@ -380,9 +380,6 @@ class Qwen2_5_VLTextModel(TextEncoder): if weight_name not in name: continue name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - # if name.endswith(".bias") and name not in params_dict: - # continue if name not in params_dict: continue param = params_dict[name] @@ -390,9 +387,6 @@ class Qwen2_5_VLTextModel(TextEncoder): weight_loader(param, loaded_weight, shard_id) break else: - # Skip loading extra bias for GPTQ models. - # if name.endswith(".bias") and name not in params_dict: - # continue if name not in params_dict: continue diff --git a/fastvideo/models/encoders/qwen2_5_vl_custom.py b/fastvideo/models/encoders/qwen2_5_vl_custom.py index e0be36be6..6f9f03722 100644 --- a/fastvideo/models/encoders/qwen2_5_vl_custom.py +++ b/fastvideo/models/encoders/qwen2_5_vl_custom.py @@ -579,9 +579,6 @@ class Qwen2_5_VLRotaryEmbedding(nn.Module): # self.register_buffer("inv_freq", inv_freq, persistent=False) # self.max_seq_len_cached = seq_len - # if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset - # self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) - # self.max_seq_len_cached = self.original_max_seq_len @torch.no_grad() def forward(self, x, position_ids): @@ -753,7 +750,6 @@ class Qwen2_5_VLAttention(nn.Module): } # Specific to RoPE models key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) - # repeat k/v heads if n_kv_heads < n_heads key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) @@ -842,7 +838,6 @@ class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): } # Specific to RoPE models key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) - # repeat k/v heads if n_kv_heads < n_heads key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) dropout_rate = 0.0 if not self.training else self.attention_dropout @@ -854,7 +849,6 @@ class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): if input_dtype == torch.float32: if torch.is_autocast_enabled(): target_dtype = torch.get_autocast_gpu_dtype() - # Handle the case where the model is quantized elif hasattr(self.config, "_pre_quantization_dtype"): target_dtype = self.config._pre_quantization_dtype else: @@ -1196,7 +1190,6 @@ class Qwen2_5_VLModel(nn.Module): hidden_states = inputs_embeds - # create position embeddings to be shared across the decoder layers position_embeddings = self.rotary_emb(hidden_states, position_ids) # decoder layers @@ -1260,7 +1253,6 @@ class Qwen2_5_VLModel(nn.Module): hidden_states = self.norm(hidden_states) - # add hidden states from the last decoder layer if output_hidden_states: all_hidden_states += (hidden_states, ) diff --git a/fastvideo/models/encoders/reason1.py b/fastvideo/models/encoders/reason1.py index c8c40ee3e..f321d67f6 100644 --- a/fastvideo/models/encoders/reason1.py +++ b/fastvideo/models/encoders/reason1.py @@ -192,7 +192,6 @@ class Reason1TextEncoder(TextEncoder): skipped_weights["decoder"] += 1 continue - # Handle stacked params mapping (for quantized models) for param_name, weight_name, shard_id in self.config.arch_config.stacked_params_mapping: if weight_name not in name: continue diff --git a/fastvideo/models/encoders/siglip.py b/fastvideo/models/encoders/siglip.py index 0ac5cf506..abea7410c 100644 --- a/fastvideo/models/encoders/siglip.py +++ b/fastvideo/models/encoders/siglip.py @@ -312,7 +312,6 @@ class SiglipVisionTransformer(nn.Module): if not return_all_hidden_states: encoder_outputs = encoder_outputs[0] - # Apply post-layernorm if self.post_layernorm is not None: if isinstance(encoder_outputs, list): encoder_outputs[-1] = self.post_layernorm(encoder_outputs[-1]) @@ -359,11 +358,9 @@ class SiglipVisionModel(ImageEncoder): layer_count = len(self.vision_model.encoder.layers) for name, loaded_weight in weights: - # Skip projection layers if any if name.startswith("visual_projection"): continue - # Skip head if any if "head" in name: continue @@ -377,7 +374,6 @@ class SiglipVisionModel(ImageEncoder): if layer_idx >= layer_count: continue - # Handle QKV projection weight mapping for (param_name, weight_name, shard_id) in self.config.arch_config.stacked_params_mapping: if weight_name not in name: continue diff --git a/fastvideo/models/encoders/t5.py b/fastvideo/models/encoders/t5.py index e85e130dc..f20a5a6b0 100644 --- a/fastvideo/models/encoders/t5.py +++ b/fastvideo/models/encoders/t5.py @@ -161,7 +161,6 @@ class T5Attention(nn.Module): self.inner_dim = self.n_heads * self.key_value_proj_dim # No GQA in t5. - # self.n_kv_heads = self.n_heads self.qkv_proj = QKVParallelLinear( self.d_model, @@ -225,7 +224,6 @@ class T5Attention(nn.Module): relative_position = torch.abs(relative_position) else: relative_position = -torch.min(relative_position, torch.zeros_like(relative_position)) - # now relative_position is in the range [0, inf) # half of the buckets are for exact increments in positions max_exact = num_buckets // 2 @@ -420,7 +418,6 @@ class T5Block(nn.Module): if self.is_decoder: hidden_states = self.layer[1](hidden_states=hidden_states, attn_metadata=attn_metadata) - # Apply Feed Forward layer hidden_states = self.layer[2](hidden_states) else: hidden_states = self.layer[1](hidden_states) diff --git a/fastvideo/models/encoders/vision.py b/fastvideo/models/encoders/vision.py index 603033262..b5075be97 100644 --- a/fastvideo/models/encoders/vision.py +++ b/fastvideo/models/encoders/vision.py @@ -83,7 +83,6 @@ def resolve_visual_encoder_outputs( for layer_idx in feature_sample_layers ] - # Apply post-norm on the final hidden state if we are using it uses_last_layer = feature_sample_layers[-1] in (len(hs_pool) - 1, -1) if post_layer_norm is not None and uses_last_layer: hs_pool[-1] = post_layer_norm(encoder_outputs) diff --git a/fastvideo/models/hf_transformer_utils.py b/fastvideo/models/hf_transformer_utils.py index 5f26179d9..cf06471c5 100644 --- a/fastvideo/models/hf_transformer_utils.py +++ b/fastvideo/models/hf_transformer_utils.py @@ -93,12 +93,10 @@ def get_diffusers_config(model: str, ) -> dict[str, Any]: config_file = os.path.join(model, config_name) if os.path.exists(config_file): try: - # Load the config directly from the file with open(config_file) as f: config_dict: dict[str, Any] = json.load(f) if "_diffusers_version" in config_dict: config_dict.pop("_diffusers_version") - # TODO(will): apply any overrides from inference args return config_dict except Exception as e: raise RuntimeError(f"Failed to load diffusers config from {config_file}: {e}") from e diff --git a/fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py b/fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py index 080f91056..703a2e939 100644 --- a/fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py +++ b/fastvideo/models/loader/benchmarks/benchmark_weight_loading_comparison.py @@ -163,7 +163,6 @@ def main(): args.repeats, ) - # Clear GPU memory between benchmarks torch.cuda.empty_cache() dist.barrier() @@ -176,7 +175,6 @@ def main(): args.repeats, ) - # Clear GPU memory between benchmarks torch.cuda.empty_cache() dist.barrier() diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 06c6bc693..ffd79b44e 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -181,8 +181,6 @@ class TextEncoderLoader(ComponentLoader): """Prepare weights for the model. If the model is not local, it will be downloaded.""" - # model_name_or_path = (self._maybe_download_from_modelscope( - # model_name_or_path, revision) or model_name_or_path) is_local = os.path.isdir(model_name_or_path) assert is_local, "Model path must be a local directory" @@ -231,7 +229,6 @@ class TextEncoderLoader(ComponentLoader): if self.counter_before_loading_weights == 0.0: self.counter_before_loading_weights = time.perf_counter() - # Apply the prefix. return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator) def _get_all_weights( @@ -257,12 +254,6 @@ class TextEncoderLoader(ComponentLoader): def load(self, model_path: str, fastvideo_args: FastVideoArgs): """Load the text encoders based on the model path, and inference args.""" - # model_config: PretrainedConfig = get_hf_config( - # model=model_path, - # trust_remote_code=fastvideo_args.trust_remote_code, - # revision=fastvideo_args.revision, - # model_override_args=None, - # ) model_config = get_diffusers_config(model=model_path) model_config.pop("_name_or_path", None) model_config.pop("transformers_version", None) @@ -362,7 +353,6 @@ class TextEncoderLoader(ComponentLoader): encoder_precision = encoder_precisions[idx] target_device = get_local_torch_device() - # TODO(will): add support for other dtypes return self.load_model( model_path, encoder_config, @@ -557,12 +547,6 @@ class ImageEncoderLoader(TextEncoderLoader): def load(self, model_path: str, fastvideo_args: FastVideoArgs): """Load the text encoders based on the model path, and inference args.""" - # model_config: PretrainedConfig = get_hf_config( - # model=model_path, - # trust_remote_code=fastvideo_args.trust_remote_code, - # revision=fastvideo_args.revision, - # model_override_args=None, - # ) with open(os.path.join(model_path, "config.json")) as f: model_config = json.load(f) model_config.pop("_name_or_path", None) @@ -610,7 +594,6 @@ class ImageEncoderLoader(TextEncoderLoader): target_device = (torch.device("mps") if current_platform.is_mps() else torch.device("cpu")) else: target_device = get_local_torch_device() - # TODO(will): add support for other dtypes return self.load_model( model_path, encoder_config, @@ -768,8 +751,6 @@ class TokenizerLoader(ComponentLoader): tokenizer = AutoTokenizer.from_pretrained( resolved_model_path, # "/tokenizer" # in v0, this was same string as encoder_name "ClipTextModel" - # TODO(will): pass these tokenizer kwargs from inference args? Maybe - # other method of config? local_files_only=os.path.isdir(resolved_model_path), ) except (OSError, ValueError): @@ -1301,7 +1282,6 @@ class TransformerLoader(ComponentLoader): default_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision] - # Load the model using FSDP loader logger.info("Loading model from %s, default_dtype: %s", cls_name, default_dtype) assert fastvideo_args.hsdp_shard_dim is not None # Cosmos2.5 checkpoints can include extra entries not present in the @@ -1356,7 +1336,6 @@ class TransformerLoader(ComponentLoader): pin_cpu_memory=fastvideo_args.pin_cpu_memory, fsdp_inference=fastvideo_args.use_fsdp_inference, device_mesh=device_mesh, - # TODO(will): make these configurable default_dtype=default_dtype, param_dtype=torch.bfloat16, reduce_dtype=torch.float32, @@ -1396,7 +1375,6 @@ class TransformerLoader(ComponentLoader): model = model.to(get_local_torch_device()) if fastvideo_args.inference_mode and fastvideo_args.dit_layerwise_offload: - # Check if model has nn.ModuleList for layerwise offload compatibility has_module_list = any(isinstance(m, nn.ModuleList) for m in model.children()) if has_module_list: enable_layerwise_offload(model) @@ -1664,7 +1642,6 @@ class PipelineComponentLoader: ) module_name = module_name.rsplit("/", 1)[-1] - # Get the appropriate loader for this module type loader = ComponentLoader.for_module_type(module_name, transformers_or_diffusers) # Resolve this component's attention backend ONCE, here, from the diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index ca8435cb6..e76bfa5a1 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -233,7 +233,6 @@ def _maybe_quantize_model(model: nn.Module, *, defer_weight_conversion_until_lor qat_train_skipped) -# TODO(PY): move this to utils elsewhere @contextlib.contextmanager def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]: """ @@ -317,8 +316,6 @@ def _validate_fsdp_inference_quantization(init_params: dict[str, Any], fsdp_infe if isinstance(quant_config, NVFP4QATTrainConfig): return - # TODO: (David) Currently reject every FSDP inference quantization config except NVFP4QATTrainConfig. - # Support FSDP inference with precomputed quantized weights. raise NotImplementedError( "FSDP inference supports unquantized transformers and NVFP4QATTrainConfig only; " f"got {type(quant_config).__name__}.") @@ -959,13 +956,11 @@ def shard_model( # Count all parameters param_count = sum(p.numel() for p in m.parameters(recurse=True)) - # Skip small modules if param_count < min_params: logger.info("Skipping module %s (%.2fM params < %.2fM threshold)", n, param_count / 1e6, min_params / 1e6) continue - # Shard this module logger.info("Sharding module %s (%.2fM params)", n, param_count / 1e6) module_kwargs = fsdp_kwargs local_ignored_params = ignored_params_by_module[id(m)] @@ -974,7 +969,6 @@ def shard_model( fully_shard(m, **module_kwargs) num_layers_sharded += 1 else: - # Shard all modules matching conditions for n, m in reversed(named_modules): if any([shard_condition(n, m) for shard_condition in fsdp_shard_conditions]): if id(m) in fp32_group_ids: @@ -1005,7 +999,6 @@ def shard_model( fully_shard(module, **fp32_kwargs) logger.info("Sharded FP32 compute modules: %s", [name for name, _ in fp32_groups]) - # Finally shard the entire model to account for any stragglers root_kwargs = fsdp_kwargs if ignored_params: root_kwargs = {**fsdp_kwargs, "ignored_params": ignored_params} @@ -1018,7 +1011,6 @@ def _drop_state_dict_parameter(model: nn.Module, param_name: str) -> None: module.register_parameter(attr_name, None) -# TODO(PY): device mesh for cfg parallel def load_model_from_full_model_state_dict( model: FSDPModule | torch.nn.Module, full_sd_iterator: Generator[tuple[str, torch.Tensor], None, None], diff --git a/fastvideo/models/loader/weight_utils.py b/fastvideo/models/loader/weight_utils.py index 1361ff179..76ae4aca3 100644 --- a/fastvideo/models/loader/weight_utils.py +++ b/fastvideo/models/loader/weight_utils.py @@ -198,8 +198,6 @@ def safetensors_weights_iterator(hf_weights_files: list[str], shape = sl.get_shape() dtype = SAFETENSORS_TO_TORCH_DTYPE[sl.get_dtype()] param = torch.empty(shape, device=device, dtype=dtype) - # broadcast to local ranks - # TODO(Wenxuan): scatter instead of broadcast if node_group.world_size > 1: group = node_group.device_group if async_broadcast: diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 897e3f276..3d890c5d2 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -319,7 +319,6 @@ def _run_in_subprocess(fn: Callable[[], _T]) -> _T: try: returned.check_returncode() except Exception as e: - # wrap raised exception to provide more information raise RuntimeError(f"Error raised in subprocess:\n" f"{returned.stderr.decode()}") from e diff --git a/fastvideo/models/schedulers/base.py b/fastvideo/models/schedulers/base.py index 758cce334..f46654550 100644 --- a/fastvideo/models/schedulers/base.py +++ b/fastvideo/models/schedulers/base.py @@ -11,7 +11,6 @@ class BaseScheduler(ABC): num_train_timesteps: int def __init__(self, *args, **kwargs) -> None: - # Check if subclass has defined all required properties required_attributes = ['timesteps', 'order', 'num_train_timesteps'] for attr in required_attributes: diff --git a/fastvideo/models/schedulers/scheduling_flow_map_euler_discrete.py b/fastvideo/models/schedulers/scheduling_flow_map_euler_discrete.py index 2108fc74c..361b0061b 100644 --- a/fastvideo/models/schedulers/scheduling_flow_map_euler_discrete.py +++ b/fastvideo/models/schedulers/scheduling_flow_map_euler_discrete.py @@ -52,8 +52,6 @@ class FlowMapEulerDiscreteScheduler(BaseScheduler): self.sigmas: torch.Tensor = torch.empty(0) super().__init__() - # ------------------------------------------------------------------ - # BaseScheduler abstract surface def set_shift(self, shift: float) -> None: self.shift = float(shift) @@ -67,7 +65,6 @@ class FlowMapEulerDiscreteScheduler(BaseScheduler): del timestep return sample - # ------------------------------------------------------------------ # Public helpers used by the AnyFlow pretrain method. def apply_shift( @@ -120,7 +117,6 @@ class FlowMapEulerDiscreteScheduler(BaseScheduler): denom = w.sum().clamp_min(1e-8) return w * (float(self.num_train_timesteps) / denom) - # ------------------------------------------------------------------ def set_timesteps( self, diff --git a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py index 7fad5b97f..2a559dfce 100644 --- a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -13,11 +13,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ============================================================================== -# # Modified from diffusers==0.29.2 -# -# ============================================================================== import math from dataclasses import dataclass from typing import Any @@ -534,7 +530,6 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler """Constructs the noise schedule of Karras et al. (2022).""" # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers sigma_min = self.config.sigma_min if hasattr(self.config, "sigma_min") else None sigma_max = self.config.sigma_max if hasattr(self.config, "sigma_max") else None @@ -554,7 +549,6 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler """Constructs an exponential noise schedule.""" # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers sigma_min = self.config.sigma_min if hasattr(self.config, "sigma_min") else None sigma_max = self.config.sigma_max if hasattr(self.config, "sigma_max") else None @@ -574,7 +568,6 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers sigma_min = self.config.sigma_min if hasattr(self.config, "sigma_min") else None sigma_max = self.config.sigma_max if hasattr(self.config, "sigma_max") else None diff --git a/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py b/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py index b03cb88a2..feeaa20b3 100644 --- a/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py +++ b/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py @@ -750,7 +750,6 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): noise: torch.Tensor, timesteps: torch.IntTensor, ) -> torch.Tensor: - # Make sure sigmas and timesteps have the same device and dtype as original_samples sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype) if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): # mps does not support float64 diff --git a/fastvideo/models/schedulers/scheduling_unipc_multistep.py b/fastvideo/models/schedulers/scheduling_unipc_multistep.py index 0b1e47f88..dff0eb33a 100644 --- a/fastvideo/models/schedulers/scheduling_unipc_multistep.py +++ b/fastvideo/models/schedulers/scheduling_unipc_multistep.py @@ -14,11 +14,7 @@ # DISCLAIMER: check https://huggingface.co/papers/2302.04867 and https://github.com/wl-zhao/UniPC for more info # The codebase is modified based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py -# ============================================================================== -# # Modified from diffusers==0.35.0.dev0 -# -# ============================================================================== import math from typing import List, Optional, Tuple, Union @@ -94,22 +90,18 @@ def rescale_zero_terminal_snr(betas): Returns: `torch.Tensor`: rescaled betas with zero terminal SNR """ - # Convert betas to alphas_bar_sqrt alphas = 1.0 - betas alphas_cumprod = torch.cumprod(alphas, dim=0) alphas_bar_sqrt = alphas_cumprod.sqrt() - # Store old values. alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() # Shift so the last timestep is zero. alphas_bar_sqrt -= alphas_bar_sqrt_T - # Scale so the first timestep is back to the old value. alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) - # Convert alphas_bar_sqrt to betas alphas_bar = alphas_bar_sqrt**2 # Revert sqrt alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod alphas = torch.cat([alphas_bar[0:1], alphas]) @@ -474,13 +466,10 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t def _sigma_to_t(self, sigma, log_sigmas): - # get log sigma log_sigma = np.log(np.maximum(sigma, 1e-10)) - # get distribution dists = log_sigma - log_sigmas[:, np.newaxis] - # get sigmas range low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) high_idx = low_idx + 1 @@ -512,7 +501,6 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): """Constructs the noise schedule of Karras et al. (2022).""" # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers sigma_min = self.config.sigma_min if hasattr(self.config, "sigma_min") else None sigma_max = self.config.sigma_max if hasattr(self.config, "sigma_max") else None @@ -532,7 +520,6 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): """Constructs an exponential noise schedule.""" # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers sigma_min = self.config.sigma_min if hasattr(self.config, "sigma_min") else None sigma_max = self.config.sigma_max if hasattr(self.config, "sigma_max") else None @@ -552,7 +539,6 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" # Hack to make sure that other schedulers which copy this function don't break - # TODO: Add this logic to the other schedulers sigma_min = self.config.sigma_min if hasattr(self.config, "sigma_min") else None sigma_max = self.config.sigma_max if hasattr(self.config, "sigma_max") else None diff --git a/fastvideo/models/utils.py b/fastvideo/models/utils.py index c83afc134..b43f73361 100644 --- a/fastvideo/models/utils.py +++ b/fastvideo/models/utils.py @@ -6,7 +6,6 @@ from typing import Any import torch -# TODO(PY): move it elsewhere def auto_attributes(init_func): """ Decorator that automatically adds all initialization arguments as object attributes. @@ -22,30 +21,25 @@ def auto_attributes(init_func): """ def wrapper(self, *args, **kwargs): - # Get the function signature import inspect signature = inspect.signature(init_func) parameters = signature.parameters - # Get parameter names (excluding 'self') param_names = list(parameters.keys())[1:] # Bind arguments to parameters bound_args = signature.bind(self, *args, **kwargs) bound_args.apply_defaults() - # Create config object if it doesn't exist if not hasattr(self, 'config'): self.config = type('Config', (), {})() - # Set attributes on self and self.config for name in param_names: if name in bound_args.arguments: value = bound_args.arguments[name] setattr(self, name, value) setattr(self.config, name, value) - # Call the original __init__ function return init_func(self, *args, **kwargs) return wrapper diff --git a/fastvideo/models/vaes/common.py b/fastvideo/models/vaes/common.py index 6cb85c982..ca7741297 100644 --- a/fastvideo/models/vaes/common.py +++ b/fastvideo/models/vaes/common.py @@ -124,8 +124,6 @@ class ParallelTiledVAE(ABC): The latent representation of the encoded videos. """ _, _, _, height, width = x.shape - # latent_height = height // self.spatial_compression_ratio - # latent_width = width // self.spatial_compression_ratio tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio @@ -166,7 +164,6 @@ class ParallelTiledVAE(ABC): world_size, rank = get_sp_world_size(), get_sp_parallel_rank() B, C, T, H, W = z.shape - # Calculate parameters tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio @@ -177,45 +174,37 @@ class ParallelTiledVAE(ABC): blend_height = self.tile_sample_min_height - self.tile_sample_stride_height blend_width = self.tile_sample_min_width - self.tile_sample_stride_width - # Calculate tile dimensions num_t_tiles = (T + tile_latent_stride_num_frames - 1) // tile_latent_stride_num_frames num_h_tiles = (H + tile_latent_stride_height - 1) // tile_latent_stride_height num_w_tiles = (W + tile_latent_stride_width - 1) // tile_latent_stride_width total_spatial_tiles = num_h_tiles * num_w_tiles total_tiles = num_t_tiles * total_spatial_tiles - # Calculate tiles per rank and padding tiles_per_rank = (total_tiles + world_size - 1) // world_size start_tile_idx = rank * tiles_per_rank end_tile_idx = min((rank + 1) * tiles_per_rank, total_tiles) local_results = [] local_dim_metadata = [] - # Process assigned tiles for local_idx, global_idx in enumerate(range(start_tile_idx, end_tile_idx)): t_idx = global_idx // total_spatial_tiles spatial_idx = global_idx % total_spatial_tiles h_idx = spatial_idx // num_w_tiles w_idx = spatial_idx % num_w_tiles - # Calculate positions t_start = t_idx * tile_latent_stride_num_frames h_start = h_idx * tile_latent_stride_height w_start = w_idx * tile_latent_stride_width - # Extract and process tile tile = z[:, :, t_start:t_start + tile_latent_min_num_frames + 1, h_start:h_start + tile_latent_min_height, w_start:w_start + tile_latent_min_width] - # Process tile tile = self._decode(tile) if t_start > 0: tile = tile[:, :, 1:, :, :] - # Store metadata shape = tile.shape - # Store decoded data (flattened) decoded_flat = tile.reshape(-1) local_results.append(decoded_flat) local_dim_metadata.append(shape) @@ -231,15 +220,12 @@ class ParallelTiledVAE(ABC): padded_results[:results.size(0)] = results del results - # Gather all results gathered_dim_metadata = [None] * world_size gathered_results = torch.zeros_like(padded_results).repeat( world_size, *[1] * len(padded_results.shape)).contiguous( ) # use contiguous to make sure it won't copy data in the following operations - # TODO (PY): use fastvideo distributed methods dist.all_gather_into_tensor(gathered_results, padded_results) dist.all_gather_object(gathered_dim_metadata, local_dim_metadata) - # Process gathered results data: list = [[[[] for _ in range(num_w_tiles)] for _ in range(num_h_tiles)] for _ in range(num_t_tiles)] for current_data, global_idx in self._parallel_data_generator(gathered_results, gathered_dim_metadata): t_idx = global_idx // total_spatial_tiles @@ -290,8 +276,6 @@ class ParallelTiledVAE(ABC): """ _, _, _, height, width = z.shape - # sample_height = height * self.spatial_compression_ratio - # sample_width = width * self.spatial_compression_ratio tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio @@ -317,7 +301,6 @@ class ParallelTiledVAE(ABC): def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: _, _, num_frames, height, width = x.shape - # tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio row = [] @@ -448,7 +431,6 @@ class DiagonalGaussianDistribution: dtype=self.parameters.dtype) def sample(self, generator: torch.Generator | None = None) -> torch.Tensor: - # make sure sample is on the same device as the parameters and has same dtype sample = randn_tensor( self.mean.shape, generator=generator, diff --git a/fastvideo/models/vaes/flux2vae.py b/fastvideo/models/vaes/flux2vae.py index c1325deec..ab0247875 100644 --- a/fastvideo/models/vaes/flux2vae.py +++ b/fastvideo/models/vaes/flux2vae.py @@ -64,7 +64,6 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): batch_norm_momentum: float = arch_config.batch_norm_momentum patch_size: Tuple[int, int] = arch_config.patch_size - # pass init params to Encoder self.encoder = Encoder( in_channels=in_channels, out_channels=latent_channels, @@ -77,7 +76,6 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): mid_block_add_attention=mid_block_add_attention, ) - # pass init params to Decoder self.decoder = Decoder( in_channels=latent_channels, out_channels=out_channels, @@ -117,7 +115,6 @@ class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE): `dict` of attention processors: A dictionary containing all attention processors used in the model with indexed by its weight name. """ - # set recursively processors = {} def fn_recursive_add_processors( diff --git a/fastvideo/models/vaes/hunyuanvae.py b/fastvideo/models/vaes/hunyuanvae.py index f6520c794..7d4401d5d 100644 --- a/fastvideo/models/vaes/hunyuanvae.py +++ b/fastvideo/models/vaes/hunyuanvae.py @@ -53,7 +53,6 @@ class HunyuanVAEAttention(nn.Module): inner_dim = heads * dim_head - # Define the projection layers self.to_q = nn.Linear(in_channels, inner_dim, bias=bias) self.to_k = nn.Linear(in_channels, inner_dim, bias=bias) self.to_v = nn.Linear(in_channels, inner_dim, bias=bias) @@ -696,8 +695,6 @@ class AutoencoderKLHunyuanVideo(nn.Module, ParallelTiledVAE): nn.Module.__init__(self) ParallelTiledVAE.__init__(self, config) - # TODO(will): only pass in config. We do this by manually defining a - # config for hunyuan vae self.block_out_channels = config.block_out_channels if config.load_encoder: diff --git a/fastvideo/models/vaes/hyworldvae.py b/fastvideo/models/vaes/hyworldvae.py index 4e95225bf..c0626ad1f 100644 --- a/fastvideo/models/vaes/hyworldvae.py +++ b/fastvideo/models/vaes/hyworldvae.py @@ -135,7 +135,6 @@ class HYWorldUpsample(nn.Module): # First chunk: only spatial upsample h = rearrange(h, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2) h = h[:, :h.shape[1] // 2] - # Compute the shortcut part shortcut = rearrange(x, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2) shortcut = shortcut.repeat_interleave(repeats=self.repeats // 2, dim=1) elif feat_cache is None and x.shape[2] > 1: @@ -199,7 +198,6 @@ class HYWorldDownsample(nn.Module): """ r1 = 2 if self.add_temporal_downsample else 1 - # Apply conv with caching if feat_cache is not None and feat_idx is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() @@ -219,7 +217,6 @@ class HYWorldDownsample(nn.Module): # Single frame: only spatial downsample h = rearrange(h, "b c f (h r2) (w r3) -> b (r2 r3 c) f h w", r2=2, r3=2) h = torch.cat([h, h], dim=1) - # Compute the shortcut part shortcut = rearrange(x, "b c f (h r2) (w r3) -> b (r2 r3 c) f h w", r2=2, r3=2) B, C, T, H, W = shortcut.shape shortcut = shortcut.view(B, h.shape[1], self.group_size // 2, T, H, W).mean(dim=2) @@ -281,7 +278,6 @@ class HYWorldResnetBlock(nn.Module): hidden_states = self.norm1(hidden_states) hidden_states = self.nonlinearity(hidden_states) - # apply the feature cacheing mechanism if feat_cache is not None and feat_idx is not None: # Retrieve the current layer index. idx = feat_idx[0] @@ -299,10 +295,8 @@ class HYWorldResnetBlock(nn.Module): dim=2, ) - # Apply the convolution layer using `hidden_states`and the *previous* cached state. hidden_states = self.conv1(hidden_states, feat_cache[idx]) - # Update the cache for this layer with the newly prepared context ('cache_x'). feat_cache[idx] = cache_x # Increment the global layer index. @@ -812,7 +806,6 @@ class AutoencoderKLHYWorld(nn.Module, ParallelTiledVAE): upsample_match_channel=config.upsample_match_channel, ) - # TODO: Add spatial tiling. self.use_tiling = False # The minimal tile height and width for spatial tiling to be used @@ -914,7 +907,6 @@ class AutoencoderKLHYWorld(nn.Module, ParallelTiledVAE): self.clear_cache() - # Process one frame at a time with caching for i in range(num_frame): self._conv_idx = [0] if i == 0: diff --git a/fastvideo/models/vaes/lingbotworld2_wanvae.py b/fastvideo/models/vaes/lingbotworld2_wanvae.py index 979fd2410..3fe8ff1e6 100644 --- a/fastvideo/models/vaes/lingbotworld2_wanvae.py +++ b/fastvideo/models/vaes/lingbotworld2_wanvae.py @@ -128,9 +128,6 @@ class Resample(nn.Module): else: cache_x = x[:, :, -1:, :, :].clone() - # if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx]!='Rep': - # # cache last frame of last two chunk - # cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) feat_cache[idx] = cache_x @@ -154,7 +151,6 @@ class Resample(nn.Module): nn.init.zeros_(conv_weight) c1, c2, t, h, w = conv_weight.size() init_matrix = torch.eye(c1 // 2, c2) - #init_matrix = repeat(init_matrix, 'o ... -> (o 2) ...').permute(1,0,2).contiguous().reshape(c1,c2) conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix conv.weight.data.copy_(conv_weight) @@ -215,10 +211,8 @@ class AttentionBlock(nn.Module): b, c, t, h, w = x.size() x = rearrange(x, 'b c t h w -> (b t) c h w') x = self.norm(x) - # compute query, key, value q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1) - # apply attention x = F.scaled_dot_product_attention( q, k, @@ -254,7 +248,6 @@ class Encoder3d(nn.Module): dims = [dim * u for u in [1] + dim_mult] scale = 1.0 - # init block self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) # downsample blocks @@ -348,7 +341,6 @@ class Decoder3d(nn.Module): dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] scale = 1.0 / 2**(len(dim_mult) - 2) - # init block self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) # middle blocks @@ -539,11 +531,9 @@ def _video_vae(pretrained_path=None, z_dim=None, device='cpu', **kwargs): dropout=0.0) cfg.update(**kwargs) - # init model with torch.device('meta'): model = WanVAE_(**cfg) - # load checkpoint logging.info(f'loading {pretrained_path}') model.load_state_dict(torch.load(pretrained_path, map_location=device), assign=True) @@ -568,7 +558,6 @@ class Wan2_1_VAE: self.std = torch.tensor(std, dtype=dtype, device=device) self.scale = [self.mean, 1.0 / self.std] - # init model self.model = _video_vae( pretrained_path=vae_pth, z_dim=z_dim, diff --git a/fastvideo/models/vaes/ltx2vae.py b/fastvideo/models/vaes/ltx2vae.py index af347064f..a90d5a096 100644 --- a/fastvideo/models/vaes/ltx2vae.py +++ b/fastvideo/models/vaes/ltx2vae.py @@ -20,9 +20,6 @@ import fastvideo.envs as envs from fastvideo.models.vaes.common import DiagonalGaussianDistribution -# ============================================================================= -# Enums -# ============================================================================= class NormLayerType(Enum): @@ -44,9 +41,6 @@ class PaddingModeType(Enum): CIRCULAR = "circular" -# ============================================================================= -# Tiling Data Structures -# ============================================================================= @dataclass(frozen=True) @@ -183,9 +177,6 @@ class Tile(NamedTuple): return combined_mask -# ============================================================================= -# Tiling Helper Functions -# ============================================================================= def compute_trapezoidal_mask_1d( @@ -349,9 +340,6 @@ def create_tiles( return create_tiles_from_intervals_and_mappers(latent_intervals, mappers) -# ============================================================================= -# Utility Functions -# ============================================================================= def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: @@ -409,9 +397,6 @@ def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> to return x -# ============================================================================= -# Normalization Layers -# ============================================================================= class PixelNorm(nn.Module): @@ -430,9 +415,6 @@ class PixelNorm(nn.Module): return x / torch.sqrt(mean_sq + self.eps) -# ============================================================================= -# Per-Channel Statistics for Latent Normalization -# ============================================================================= class PerChannelStatistics(nn.Module): @@ -460,9 +442,6 @@ class PerChannelStatistics(nn.Module): return (x - mean) / std -# ============================================================================= -# Causal 3D Convolution -# ============================================================================= class CausalConv3d(nn.Module): @@ -591,9 +570,6 @@ def make_linear_nd( raise ValueError(f"unsupported dimensions: {dims}") -# ============================================================================= -# ResNet Blocks -# ============================================================================= class ResnetBlock3D(nn.Module): @@ -814,9 +790,6 @@ class UNetMidBlock3D(nn.Module): return hidden_states -# ============================================================================= -# Timestep Embedding (simplified version for decoder conditioning) -# ============================================================================= class PixArtAlphaCombinedTimestepSizeEmbeddings(nn.Module): @@ -870,9 +843,6 @@ class TimestepEmbedding(nn.Module): return sample -# ============================================================================= -# Sampling Blocks (Downsampling / Upsampling) -# ============================================================================= class SpaceToDepthDownsample(nn.Module): @@ -982,9 +952,6 @@ class DepthToSpaceUpsample(nn.Module): return x -# ============================================================================= -# Block Factory Functions -# ============================================================================= def _make_encoder_block( @@ -1178,9 +1145,6 @@ def _make_decoder_block( return block, out_channels -# ============================================================================= -# Video Encoder -# ============================================================================= class VideoEncoder(nn.Module): @@ -1304,9 +1268,6 @@ class VideoEncoder(nn.Module): return self.per_channel_statistics.normalize(means) -# ============================================================================= -# Video Decoder -# ============================================================================= class VideoDecoder(nn.Module): @@ -1475,9 +1436,6 @@ class VideoDecoder(nn.Module): return sample -# ============================================================================= -# Configurators (for loading from config dict) -# ============================================================================= class VideoEncoderConfigurator: @@ -1538,9 +1496,6 @@ class VideoDecoderConfigurator: ) -# ============================================================================= -# Public API (Wrapper Classes) -# ============================================================================= class LTX2VideoEncoder(nn.Module): @@ -1797,7 +1752,6 @@ class LTX2CausalVideoAutoencoder(nn.Module): yield_len = curr_temporal_slice.start - previous_temporal_slice.start yield (previous_chunk / previous_weights)[:, :, :yield_len, :, :] - # Update state for next iteration previous_chunk = buffer previous_weights = curr_weights previous_temporal_slice = curr_temporal_slice diff --git a/fastvideo/models/vision_utils.py b/fastvideo/models/vision_utils.py index feffe2178..0ce9cdbac 100644 --- a/fastvideo/models/vision_utils.py +++ b/fastvideo/models/vision_utils.py @@ -175,13 +175,11 @@ def _load_gif(gif_path: str) -> tuple[list[PIL.Image.Image], float | None]: original_fps = None with PIL.Image.open(gif_path) as gif: - # Extract FPS from GIF metadata if hasattr(gif, 'info') and 'duration' in gif.info: duration_ms = gif.info['duration'] if duration_ms > 0: original_fps = 1000.0 / duration_ms - # Extract all frames try: while True: pil_images.append(gif.copy()) @@ -206,7 +204,6 @@ def _load_video_with_ffmpeg(video_path: str) -> tuple[list[PIL.Image.Image], flo Raises: AttributeError: If ffmpeg is not installed """ - # Verify ffmpeg is available try: imageio.plugins.ffmpeg.get_exe() except AttributeError as e: @@ -227,7 +224,6 @@ def _load_video_with_ffmpeg(video_path: str) -> tuple[list[PIL.Image.Image], flo if isinstance(source_size, dict): original_fps = source_size.get('fps') - # Extract all frames for frame in reader: pil_images.append(PIL.Image.fromarray(frame)) @@ -291,7 +287,6 @@ def load_video( else: pil_images, original_fps = _load_video_with_ffmpeg(video_path) finally: - # Clean up temporary file if it was created if was_tempfile_created and os.path.exists(video_path): os.remove(video_path) @@ -435,7 +430,6 @@ def preprocess_reference_image_for_clip(image: PIL.Image.Image, device: torch.de Returns: Preprocessed PIL image ready for CLIP encoder """ - # Convert PIL to tensor and normalize to [-1, 1] range image_tensor = TF.to_tensor(image).sub_(0.5).div_(0.5).to(device) # Resize to CLIP's expected input size (224x224) using bicubic interpolation diff --git a/fastvideo/models/wan/causal_transformer.py b/fastvideo/models/wan/causal_transformer.py index a0a33aabf..b6c8c5b18 100644 --- a/fastvideo/models/wan/causal_transformer.py +++ b/fastvideo/models/wan/causal_transformer.py @@ -146,7 +146,6 @@ class CausalWanSelfAttention(nn.Module): value=padded_v[:, :, head_slice].transpose(2, 1), block_mask=block_mask)[:, :, :q.shape[1]].transpose(2, 1) if sp_world_size > 1: - # Gather per-rank head outputs back to the full head dimension. x = sequence_model_parallel_all_gather(x.contiguous(), dim=2) else: current_end = current_start + q.shape[1] @@ -188,12 +187,10 @@ class CausalWanSelfAttention(nn.Module): kv_cache["k"][:, local_start_index:local_end_index] = stored_key kv_cache["v"][:, local_start_index:local_end_index] = v else: - # Assign new keys/values directly up to current_end local_end_index = local_end_index_prev + current_end - global_end_index local_start_index = local_end_index - num_new_tokens kv_cache["k"] = kv_cache["k"].detach() kv_cache["v"] = kv_cache["v"].detach() - # logger.info("kv_cache['k'] is in comp graph: %s", kv_cache["k"].requires_grad or kv_cache["k"].grad_fn is not None) kv_cache["k"][:, local_start_index:local_end_index] = stored_key kv_cache["v"][:, local_start_index:local_end_index] = v key_window = kv_cache["k"][:, max(0, local_end_index - max_attention_size):local_end_index] @@ -208,7 +205,6 @@ class CausalWanSelfAttention(nn.Module): x = self.attn(roped_query[:, :, head_slice], key_window[:, :, head_slice], value_window[:, :, head_slice]) if sp_world_size > 1: - # Gather per-rank head outputs back to the full head dimension. x = sequence_model_parallel_all_gather(x.contiguous(), dim=2) if isinstance(kv_cache["global_end_index"], torch.Tensor): kv_cache["global_end_index"].fill_(current_end) @@ -274,8 +270,6 @@ class CausalWanTransformerBlock(nn.Module): elementwise_affine=True, dtype=torch.float32) - # 2. Cross-attention - # Only T2V for now self.attn2 = WanT2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps) self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, norm_type="layer", @@ -314,13 +308,11 @@ class CausalWanTransformerBlock(nn.Module): frame_seqlen = int(frame_seqlen) bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype - # assert orig_dtype != torch.float32 e = self.scale_shift_table + temb # e.shape: [batch_size, temb_seq_len, 6, inner_dim] assert e.shape == (bs, temb_seq_len, 6, self.hidden_dim) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=2) # *_msa.shape: [batch_size, temb_seq_len, 1, inner_dim] - # assert shift_msa.dtype == torch.float32 # 1. Self-attention norm_hidden_states = (self.norm1(hidden_states).unflatten(dim=1, sizes=(temb_seq_len, tokens_per_temb)) * @@ -481,7 +473,6 @@ class CausalWanTransformer3DModel(BaseDiT): else: return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx) - # return ((kv_idx < total_length) & (q_idx < total_length)) | (q_idx == kv_idx) # bidirectional mask block_mask = create_block_mask(attention_mask, B=None, @@ -495,15 +486,7 @@ class CausalWanTransformer3DModel(BaseDiT): print(f" cache a block wise causal mask with block size of {num_frame_per_block} frames") print(block_mask) - # import imageio - # import numpy as np - # from torch.nn.attention.flex_attention import create_mask - # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length + - # padded_length, KV_LEN=total_length + padded_length, device=device) - # import cv2 - # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024)) - # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask)) return block_mask @@ -611,7 +594,6 @@ class CausalWanTransformer3DModel(BaseDiT): post_patch_height = height // p_h post_patch_width = width // p_w - # Get rotary embeddings d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] if self.rope_cache_policy == "relativistic": @@ -721,7 +703,6 @@ class CausalWanTransformer3DModel(BaseDiT): post_patch_height = height // p_h post_patch_width = width // p_w - # Get rotary embeddings d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] freqs_cos, freqs_sin = get_rotary_pos_embed( @@ -855,7 +836,6 @@ class CausalWanTransformer3DModel(BaseDiT): for u, v in zip(x, grid_sizes.tolist()): u = u[:math.prod(v)].view(*v, *self.patch_size, c) u = u.permute(6, 0, 3, 1, 4, 2, 5) - # u = torch.einsum('fhwpqrc->cfphqwr', u.contiguous()) u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) out.append(u) return out diff --git a/fastvideo/models/wan/transformer.py b/fastvideo/models/wan/transformer.py index 4c9978b85..1daf0a557 100644 --- a/fastvideo/models/wan/transformer.py +++ b/fastvideo/models/wan/transformer.py @@ -15,8 +15,6 @@ from fastvideo.distributed.communication_op import (sequence_model_parallel_all_ from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidual, ScaleResidualLayerNormScaleShift) from fastvideo.layers.linear import ReplicatedLinear -# from torch.nn import RMSNorm -# TODO: RMSNorm .... from fastvideo.layers.mlp import MLP from fastvideo.layers.rotary_embedding import get_rotary_pos_embed from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, TimestepEmbedder) @@ -320,7 +318,6 @@ class WanT2VCrossAttention(WanSelfAttention): """ b, n, d = x.size(0), self.num_heads, self.head_dim - # compute query, key, value q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) if crossattn_cache is not None: @@ -337,7 +334,6 @@ class WanT2VCrossAttention(WanSelfAttention): k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) v = self.to_v(context)[0].view(b, -1, n, d) - # compute attention x = self.attn(q, k, v) if k.size(1) > 0 else torch.zeros_like(q) # output @@ -385,14 +381,12 @@ class WanI2VCrossAttention(WanSelfAttention): context = context[:, 257:] b, n, d = x.size(0), self.num_heads, self.head_dim - # compute query, key, value q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) v = self.to_v(context)[0].view(b, -1, n, d) k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view(b, -1, n, d) v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d) img_x = self.attn(q, k_img, v_img) - # compute attention x = self.attn(q, k, v) if k.size(1) > 0 else torch.zeros_like(q) # output @@ -496,7 +490,6 @@ class WanTransformerBlock(nn.Module): hidden_states = hidden_states.squeeze(1) bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype - # assert orig_dtype != torch.float32 if hidden_states.is_cuda and torch.compiler.is_compiling(): shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( @@ -668,7 +661,6 @@ class WanTransformerBlock_VSA(nn.Module): hidden_states = hidden_states.squeeze(1) bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype - # assert orig_dtype != torch.float32 if hidden_states.is_cuda and torch.compiler.is_compiling(): shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( torch.ops.fastvideo._wan_modulation_forward(self.scale_shift_table, temb)) @@ -822,7 +814,6 @@ class WanTransformer3DModel(BaseDiT): post_patch_height = height // p_h post_patch_width = width // p_w - # Get rotary embeddings d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] freqs_cos, freqs_sin = get_rotary_pos_embed((post_patch_num_frames, post_patch_height, post_patch_width), @@ -901,7 +892,6 @@ class WanTransformer3DModel(BaseDiT): hidden_states = self.norm_out(hidden_states, shift, scale) - # Gather and unpad in one operation hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1) hidden_states = self.proj_out(hidden_states) diff --git a/fastvideo/models/wan/vae.py b/fastvideo/models/wan/vae.py index e26cd8955..e750a9ce1 100644 --- a/fastvideo/models/wan/vae.py +++ b/fastvideo/models/wan/vae.py @@ -188,7 +188,6 @@ class WanCausalConv3d(nn.Conv3d): padding=padding, ) self.padding: tuple[int, int, int] - # Set up causal padding self._padding: tuple[int, ...] = (self.padding[2], self.padding[2], self.padding[1], self.padding[1], 2 * self.padding[0], 0) self.padding = (0, 0, 0) @@ -412,7 +411,6 @@ class WanResidualBlock(nn.Module): self.conv_shortcut = WanCausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() def forward(self, x): - # Apply shortcut connection h = self.conv_shortcut(x) # First normalization and activation @@ -458,7 +456,6 @@ class WanResidualBlock(nn.Module): else: x = self.conv2(x) - # Add residual connection return x + h @@ -486,13 +483,11 @@ class WanAttentionBlock(nn.Module): x = x.permute(0, 2, 1, 3, 4).reshape(batch_size * time, channels, height, width) x = self.norm(x) - # compute query, key, value qkv = self.to_qkv(x) qkv = qkv.reshape(batch_size * time, 1, channels * 3, -1) qkv = qkv.permute(0, 1, 3, 2).contiguous() q, k, v = qkv.chunk(3, dim=-1) - # apply attention x = F.scaled_dot_product_attention(q, k, v) x = x.squeeze(1).permute(0, 2, 1).reshape(batch_size * time, channels, height, width) @@ -521,7 +516,6 @@ class WanMidBlock(nn.Module): super().__init__() self.dim = dim - # Create the components resnets = [WanResidualBlock(dim, dim, dropout, non_linearity)] attentions = [] for _ in range(num_layers): @@ -536,7 +530,6 @@ class WanMidBlock(nn.Module): # First residual block x = self.resnets[0](x) - # Process through attention and residual blocks for attn, resnet in zip(self.attentions, self.resnets[1:], strict=True): if attn is not None: x = attn(x) @@ -566,7 +559,6 @@ class WanResidualDownBlock(nn.Module): in_dim = out_dim self.resnets = nn.ModuleList(resnets) - # Add the final downsample block if down_flag: mode = "downsample3d" if temperal_downsample else "downsample2d" self.downsampler = WanResample(out_dim, mode=mode) @@ -625,7 +617,6 @@ class WanEncoder3d(nn.Module): dims = [dim * u for u in [1] + dim_mult] scale = 1.0 - # init block self.conv_in = WanCausalConv3d(in_channels, dims[0], 3, padding=1) # downsample blocks @@ -748,7 +739,6 @@ class WanResidualUpBlock(nn.Module): else: self.avg_shortcut = None - # create residual blocks resnets = [] current_dim = in_dim for _ in range(num_res_blocks + 1): @@ -757,7 +747,6 @@ class WanResidualUpBlock(nn.Module): self.resnets = nn.ModuleList(resnets) - # Add upsampling layer if needed if up_flag: upsample_mode = "upsample3d" if temperal_upsample else "upsample2d" self.upsampler = WanResample(out_dim, mode=upsample_mode, upsample_out_dim=out_dim) @@ -817,9 +806,7 @@ class WanUpBlock(nn.Module): self.in_dim = in_dim self.out_dim = out_dim - # Create layers list resnets = [] - # Add residual blocks and attention if needed current_dim = in_dim for _ in range(num_res_blocks + 1): resnets.append(WanResidualBlock(current_dim, out_dim, dropout, non_linearity)) @@ -827,7 +814,6 @@ class WanUpBlock(nn.Module): self.resnets = nn.ModuleList(resnets) - # Add upsampling layer if needed self.upsamplers = None if upsample_mode is not None: self.upsamplers = nn.ModuleList([WanResample(out_dim, mode=upsample_mode)]) @@ -896,7 +882,6 @@ class WanDecoder3d(nn.Module): # dimensions dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] - # init block self.conv_in = WanCausalConv3d(z_dim, dims[0], 3, padding=1) # middle blocks @@ -919,7 +904,6 @@ class WanDecoder3d(nn.Module): elif up_flag: upsample_mode = "upsample2d" - # Create and add the upsampling block if is_residual: up_block = WanResidualUpBlock( in_dim=in_dim, diff --git a/fastvideo/pipelines/basic/gen3c/cache_3d.py b/fastvideo/pipelines/basic/gen3c/cache_3d.py index fc68e9746..71a67a696 100644 --- a/fastvideo/pipelines/basic/gen3c/cache_3d.py +++ b/fastvideo/pipelines/basic/gen3c/cache_3d.py @@ -301,11 +301,9 @@ def forward_warp( intrinsic2 = intrinsic1.clone() if world_points1 is not None: - # Use pre-computed world points assert world_points1.shape == (b, h, w, 3) trans_points1 = project_points(world_points1, transformation2, intrinsic2) else: - # Compute from depth assert depth1 is not None and transformation1 is not None assert depth1.shape == (b, 1, h, w) @@ -427,7 +425,6 @@ class Cache3DBase: if input_mask is not None: input_image = torch.cat([input_image, input_mask], dim=format_to_indices.get("C")) - # Extract dimensions B = input_shape[format_to_indices.get("B", 0)] if "B" in format_to_indices else 1 F = input_shape[format_to_indices.get("F", 0)] if "F" in format_to_indices else 1 N = input_shape[format_to_indices.get("N", 0)] if "N" in format_to_indices else 1 @@ -458,7 +455,6 @@ class Cache3DBase: self.input_image = input_image self.input_image = self.input_image.to(weight_dtype).to("cpu") - # Compute 3D world points if input_points is not None: self.input_points = input_points.reshape(B, F, N, V, H, W, 3).to("cpu") self.input_depth = None @@ -523,7 +519,6 @@ class Cache3DBase: target_intrinsics = target_intrinsics.reshape(B, F_target, 1, 3, 3).expand(B, F_target, N, 3, 3).reshape(-1, 3, 3) - # Prepare inputs first_images = rearrange( self.input_image[:, start_frame_idx:start_frame_idx + F_target].expand(B, F_target, N, V, C, H, W), "B F N V C H W -> (B F N) V C H W") @@ -534,7 +529,6 @@ class Cache3DBase: self.input_mask[:, start_frame_idx:start_frame_idx + F_target].expand(B, F_target, N, V, 1, H, W), "B F N V C H W -> (B F N) V C H W") if self.input_mask is not None else None - # Process in chunks for memory efficiency if first_images.shape[1] == 1: warp_chunk_size = 2 rendered_warp_images = [] @@ -651,7 +645,6 @@ class Cache3DBuffer(Cache3DBase): B, F, N, V, C, H, W = self.input_image.shape - # Compute new 3D points new_points = unproject_points(new_depth, new_w2c, new_intrinsics, is_depth=self.is_depth).cpu() new_image = new_image.cpu() @@ -710,7 +703,6 @@ class Cache3DBuffer(Cache3DBase): pixels = pixels.to(output_device) masks = masks.to(output_device) - # Apply noise augmentation (stronger for older buffers) if not render_depth and self.noise_aug_strength > 0: noise = torch.randn(pixels.shape, generator=self.generator, device=pixels.device, dtype=pixels.dtype) per_buffer_noise = (torch.arange(start=pixels.shape[2] - 1, end=-1, step=-1, device=pixels.device) * diff --git a/fastvideo/pipelines/basic/gen3c/depth_estimation.py b/fastvideo/pipelines/basic/gen3c/depth_estimation.py index a13f0c33f..807d4d930 100644 --- a/fastvideo/pipelines/basic/gen3c/depth_estimation.py +++ b/fastvideo/pipelines/basic/gen3c/depth_estimation.py @@ -127,7 +127,6 @@ def _predict_depth_core( resized = cv2.resize(input_image_rgb, (depth_pred_w, depth_pred_h)) img_tensor = torch.tensor(resized / 255.0, dtype=torch.float32, device=device).permute(2, 0, 1) - # Run MoGe inference moge_output = moge_model.infer(img_tensor) depth_hw = moge_output["depth"] intrinsics_norm = moge_output["intrinsics"] @@ -136,14 +135,12 @@ def _predict_depth_core( # Replace invalid depth with large value depth_hw = torch.where(mask_hw == 0, torch.tensor(1000.0, device=depth_hw.device), depth_hw) - # Convert normalized intrinsics to pixel coordinates intrinsics_pixel = intrinsics_norm.clone() intrinsics_pixel[0, 0] *= depth_pred_w # fx intrinsics_pixel[1, 1] *= depth_pred_h # fy intrinsics_pixel[0, 2] *= depth_pred_w # cx intrinsics_pixel[1, 2] *= depth_pred_h # cy - # Scale to target resolution h_scale = target_h / depth_pred_h w_scale = target_w / depth_pred_w @@ -159,7 +156,6 @@ def _predict_depth_core( img_target = F.interpolate(img_tensor.unsqueeze(0), size=(target_h, target_w), mode='bilinear', align_corners=False).squeeze(0) - # Scale intrinsics for target resolution intrinsics_target = intrinsics_pixel.clone() intrinsics_target[0, 0] *= w_scale # fx intrinsics_target[0, 2] *= w_scale # cx diff --git a/fastvideo/pipelines/basic/hunyuan/hunyuan_pipeline.py b/fastvideo/pipelines/basic/hunyuan/hunyuan_pipeline.py index e0b31ae79..5013a70d5 100644 --- a/fastvideo/pipelines/basic/hunyuan/hunyuan_pipeline.py +++ b/fastvideo/pipelines/basic/hunyuan/hunyuan_pipeline.py @@ -12,7 +12,6 @@ from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, DenoisingStage, InputValidationStage, LatentPreparationStage, TextEncodingStage, TimestepPreparationStage) -# TODO(will): move PRECISION_TO_TYPE to better place logger = init_logger(__name__) diff --git a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_2sr_pipeline.py b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_2sr_pipeline.py index 5f8a3e6cf..827f0c4f9 100644 --- a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_2sr_pipeline.py +++ b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_2sr_pipeline.py @@ -16,7 +16,6 @@ from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, Denois Hy15ImageEncodingStage, SRDenoisingStage) from fastvideo.distributed.parallel_state import get_local_torch_device -# TODO(will): move PRECISION_TO_TYPE to better place logger = init_logger(__name__) @@ -98,7 +97,6 @@ class HunyuanVideo152SRPipeline(ComposedPipelineBase): self.get_module("transformer").to(get_local_torch_device()) # Execute each stage logger.info("Running pipeline stages: %s", self._stage_name_mapping.keys()) - # logger.info("Batch: %s", batch) batch = self.input_validation_stage(batch, fastvideo_args) batch = self.prompt_encoding_stage_primary(batch, fastvideo_args) batch = self.conditioning_stage(batch, fastvideo_args) @@ -137,7 +135,6 @@ class HunyuanVideo152SRPipeline(ComposedPipelineBase): end_time = time.time() logger.info("Decoding time: %s seconds", end_time - start_time) - # Return the output return batch diff --git a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py index 1837aafe9..c55c0d1b7 100644 --- a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py +++ b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py @@ -13,7 +13,6 @@ from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, Denois LatentPreparationStage, TextEncodingStage, TimestepPreparationStage, Hy15ImageEncodingStage) -# TODO(will): move PRECISION_TO_TYPE to better place logger = init_logger(__name__) diff --git a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_pipeline.py b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_pipeline.py index 2e146dabd..86f16c321 100644 --- a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_pipeline.py +++ b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_pipeline.py @@ -13,7 +13,6 @@ from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, Denois LatentPreparationStage, TextEncodingStage, TimestepPreparationStage, Hy15ImageEncodingStage) -# TODO(will): move PRECISION_TO_TYPE to better place logger = init_logger(__name__) diff --git a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_sr_pipeline.py b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_sr_pipeline.py index a81402988..580be140e 100644 --- a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_sr_pipeline.py +++ b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_sr_pipeline.py @@ -16,7 +16,6 @@ from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, Denois Hy15ImageEncodingStage, SRDenoisingStage) from fastvideo.distributed.parallel_state import get_local_torch_device -# TODO(will): move PRECISION_TO_TYPE to better place logger = init_logger(__name__) @@ -88,7 +87,6 @@ class HunyuanVideo15SRPipeline(ComposedPipelineBase): # Execute each stage logger.info("Running pipeline stages: %s", self._stage_name_mapping.keys()) - # logger.info("Batch: %s", batch) self.get_module("transformer").to(get_local_torch_device()) batch = self.input_validation_stage(batch, fastvideo_args) batch = self.prompt_encoding_stage_primary(batch, fastvideo_args) @@ -114,7 +112,6 @@ class HunyuanVideo15SRPipeline(ComposedPipelineBase): end_time = time.time() logger.info("Decoding time: %s seconds", end_time - start_time) - # Return the output return batch diff --git a/fastvideo/pipelines/basic/lingbot_video/stages.py b/fastvideo/pipelines/basic/lingbot_video/stages.py index 03cb2e3ce..d3b9a80d1 100644 --- a/fastvideo/pipelines/basic/lingbot_video/stages.py +++ b/fastvideo/pipelines/basic/lingbot_video/stages.py @@ -293,7 +293,6 @@ class LingBotVideoDenoisingStage(PipelineStage): trajectory: list[torch.Tensor] = [] trajectory_timesteps: list[torch.Tensor] = [] for timestep in batch.timesteps: - # Stop if interrupted if getattr(self, "interrupt", False): break timestep_batch = self._transformer_timestep(timestep, transformer_dtype).expand(1).to(device) diff --git a/fastvideo/pipelines/basic/lingbotworld2/causal_fast_pipeline.py b/fastvideo/pipelines/basic/lingbotworld2/causal_fast_pipeline.py index 398562da4..7259fb00a 100644 --- a/fastvideo/pipelines/basic/lingbotworld2/causal_fast_pipeline.py +++ b/fastvideo/pipelines/basic/lingbotworld2/causal_fast_pipeline.py @@ -308,7 +308,6 @@ class LingBotWorld2CausalFastGenerationStage(PipelineStage): } x0 = current_latent for timestep_idx, timestep_value in enumerate(timesteps): - # Stop if interrupted if getattr(self, "interrupt", False): break timestep = torch.stack([timestep_value]).to(device) diff --git a/fastvideo/pipelines/basic/ltx2/stages/ltx2_audio_decoding.py b/fastvideo/pipelines/basic/ltx2/stages/ltx2_audio_decoding.py index 5eaff1f0c..f5bfb1bca 100644 --- a/fastvideo/pipelines/basic/ltx2/stages/ltx2_audio_decoding.py +++ b/fastvideo/pipelines/basic/ltx2/stages/ltx2_audio_decoding.py @@ -68,7 +68,6 @@ class LTX2AudioDecodingStage(PipelineStage): decoded_spec = self.audio_decoder(audio_latents) audio_wave = self.vocoder(decoded_spec).squeeze(0).float() - # Move to CPU for pickling across process boundary batch.extra["audio"] = audio_wave.cpu() batch.extra["audio_sample_rate"] = _resolve_audio_sample_rate(self.vocoder) return batch diff --git a/fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py b/fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py index d0c87ebdf..9eb8f2eb7 100644 --- a/fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py +++ b/fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py @@ -281,7 +281,6 @@ class LTX2DenoisingStage(PipelineStage): "falling back to no-CFG for this stage.") use_cfg = False - # Ensure text conditioning is on the same device as latents. if prompt_embeds.device != latents.device: prompt_embeds = prompt_embeds.to(latents.device) if prompt_mask is not None and prompt_mask.device != latents.device: @@ -551,7 +550,6 @@ class LTX2DenoisingStage(PipelineStage): vsa_metadata_builder = (VideoSparseAttentionMetadataBuilder() if use_vsa else None) for step_index in tqdm(range(len(sigmas) - 1)): - # Stop if interrupted if getattr(self, "interrupt", False): break sigma = sigmas[step_index] diff --git a/fastvideo/pipelines/basic/ltx2/stages/ltx2_text_encoding.py b/fastvideo/pipelines/basic/ltx2/stages/ltx2_text_encoding.py index c409e29da..fb6caa640 100644 --- a/fastvideo/pipelines/basic/ltx2/stages/ltx2_text_encoding.py +++ b/fastvideo/pipelines/basic/ltx2/stages/ltx2_text_encoding.py @@ -45,10 +45,8 @@ class LTX2TextEncodingStage(TextEncodingStage): # SP enabled: only rank 0 encodes, then broadcasts if sp_rank == 0: logger.info("[LTX2TextEncodingStage] SP rank 0: running text encoding") - # Run encoding on rank 0 result_batch = super().forward(batch, fastvideo_args) - # Build broadcast dict from batch broadcast_dict = self._build_broadcast_dict(result_batch) # Broadcast to other ranks diff --git a/fastvideo/pipelines/basic/magi_human/stages/denoising.py b/fastvideo/pipelines/basic/magi_human/stages/denoising.py index 44079d5a0..d63350da5 100644 --- a/fastvideo/pipelines/basic/magi_human/stages/denoising.py +++ b/fastvideo/pipelines/basic/magi_human/stages/denoising.py @@ -152,7 +152,6 @@ class MagiHumanDenoisingStage(PipelineStage): disable_tqdm = not getattr(fastvideo_args, "log_level_progress", True) for idx, t in enumerate(tqdm(timesteps, disable=disable_tqdm)): - # Stop if interrupted if getattr(self, "interrupt", False): break video_latent = _overwrite_first_frame(video_latent, image_latent) diff --git a/fastvideo/pipelines/basic/magi_human/stages/sr_denoising.py b/fastvideo/pipelines/basic/magi_human/stages/sr_denoising.py index 566fae08f..596194239 100644 --- a/fastvideo/pipelines/basic/magi_human/stages/sr_denoising.py +++ b/fastvideo/pipelines/basic/magi_human/stages/sr_denoising.py @@ -101,7 +101,6 @@ class MagiHumanSRDenoisingStage(PipelineStage): disable_tqdm = not getattr(fastvideo_args, "log_level_progress", True) for idx, t in enumerate(tqdm(video_scheduler.timesteps, disable=disable_tqdm)): - # Stop if interrupted if getattr(self, "interrupt", False): break video_latent = _overwrite_first_frame(video_latent, image_latent) diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py index 46b3d83d8..112902e03 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py @@ -383,7 +383,6 @@ class MiniMaxH3DenoisingStage(PipelineStage): with profiler_region("inference_denoising"), nvtx_range("minimax_h3.dit"): for index, (video_timestep, audio_timestep) in enumerate(zip(video_timesteps, audio_timesteps, strict=True)): - # Stop if interrupted if getattr(self, "interrupt", False): break unique_timesteps, timestep_indices = row_timestep_plan[index] diff --git a/fastvideo/pipelines/basic/mmaudio/stages.py b/fastvideo/pipelines/basic/mmaudio/stages.py index 339ffa0a8..6c965f21c 100644 --- a/fastvideo/pipelines/basic/mmaudio/stages.py +++ b/fastvideo/pipelines/basic/mmaudio/stages.py @@ -558,7 +558,6 @@ class MMAudioDenoisingStage(PipelineStage): empty_conditions = batch.extra["mmaudio_empty_conditions"] latents = batch.latents for index, timestep in enumerate(self.scheduler.timesteps): - # Stop if interrupted if getattr(self, "interrupt", False): break # Converted inference checkpoints carry bf16 parameters directly, diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/sampling.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/sampling.py index 017415fb5..8e244d7ca 100644 --- a/fastvideo/pipelines/basic/open_h3_regenerate_2k/sampling.py +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/sampling.py @@ -7,7 +7,6 @@ import torch REFINE_SHIFT = 6.0 DEFAULT_REFINE_STEPS = 3 -# Start noise of the two-step variant; the three-step grid starts at 0.909375. TWO_STEP_START_SIGMA = 0.7 diff --git a/fastvideo/pipelines/basic/wan/stages/causal_denoising.py b/fastvideo/pipelines/basic/wan/stages/causal_denoising.py index 4bce52006..8553eddef 100644 --- a/fastvideo/pipelines/basic/wan/stages/causal_denoising.py +++ b/fastvideo/pipelines/basic/wan/stages/causal_denoising.py @@ -148,7 +148,6 @@ class CausalDMDDenosingStage(WanCausalDenoisingBase): patch_ratio = self.transformer.config.arch_config.patch_size[ -1] * self.transformer.config.arch_config.patch_size[-2] self.frame_seq_length = latent_seq_length // patch_ratio - # TODO(will): make this a parameter once we add i2v support independent_first_frame = self.transformer.independent_first_frame if hasattr( self.transformer, 'independent_first_frame') else False # Timesteps for DMD @@ -183,11 +182,9 @@ class CausalDMDDenosingStage(WanCausalDenoisingBase): prompt_embeds = batch.prompt_embeds assert torch.isnan(prompt_embeds[0]).sum() == 0 - # Initialize or reset caches kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], dtype=target_dtype, device=latents.device) kv_cache2 = None if boundary_timestep is not None: - # Initialize the low noise kv cache kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0], dtype=target_dtype, device=latents.device) @@ -216,7 +213,6 @@ class CausalDMDDenosingStage(WanCausalDenoisingBase): block_sizes = [self.num_frames_per_block] * num_blocks start_index = 0 - # For now hardcode the first block to be 1 frame assuming the model is Wan2.2-MoE if boundary_timestep is not None: block_sizes[0] = 1 @@ -269,21 +265,18 @@ class CausalDMDDenosingStage(WanCausalDenoisingBase): video_raw_latent_shape = noise_latents_btchw.shape for i, t_cur in enumerate(timesteps): - # Stop if interrupted if getattr(self, "interrupt", False): break if boundary_timestep is not None and t_cur < boundary_timestep: current_model = self.transformer_2 else: current_model = self.transformer - # Copy for pred conversion noise_latents = noise_latents_btchw.clone() latent_model_input = current_latents.to(target_dtype) if batch.image_latent is not None and independent_first_frame and start_index == 0: latent_model_input = torch.cat([latent_model_input, batch.image_latent.to(target_dtype)], dim=2) - # Prepare inputs t_expand = t_cur.repeat(latent_model_input.shape[0]) # Attention metadata if needed @@ -310,7 +303,6 @@ class CausalDMDDenosingStage(WanCausalDenoisingBase): set_forward_context(current_timestep=i, attn_metadata=attn_metadata, forward_batch=batch): - # Run transformer; follow DMD stage pattern t_expanded_noise = t_cur * torch.ones( (latent_model_input.shape[0], 1), device=latent_model_input.device, dtype=torch.long) pred_noise_btchw = current_model( @@ -325,7 +317,6 @@ class CausalDMDDenosingStage(WanCausalDenoisingBase): **pos_cond_kwargs, ).permute(0, 2, 1, 3, 4) - # Convert pred noise to pred video with FM Euler scheduler utilities if boundary_timestep is not None and t_cur >= boundary_timestep: pred_video_btchw = pred_noise_to_x_bound( pred_noise=pred_noise_btchw.flatten(0, 1), @@ -367,7 +358,6 @@ class CausalDMDDenosingStage(WanCausalDenoisingBase): if progress_bar is not None: progress_bar.update() - # Write back and advance latents[:, :, start_index:start_index + current_num_frames, :, :] = current_latents # Re-run with context timestep to update KV cache using clean context @@ -455,7 +445,6 @@ class CausalDenoisingStage(WanCausalDenoisingBase): num_inference_steps = batch.num_inference_steps - # Initialize caches kv_cache = self._initialize_kv_cache( batch_size=b, dtype=target_dtype, @@ -502,7 +491,6 @@ class CausalDenoisingStage(WanCausalDenoisingBase): timesteps = self.scheduler.timesteps for i, t_cur in enumerate(timesteps): - # Stop if interrupted if getattr(self, "interrupt", False): break latent_model_input = current_latents.to(target_dtype) @@ -554,7 +542,6 @@ class CausalDenoisingStage(WanCausalDenoisingBase): if progress_bar is not None: progress_bar.update() - # Write denoised block back latents[ :, :, @@ -563,7 +550,6 @@ class CausalDenoisingStage(WanCausalDenoisingBase): :, ] = current_latents - # Update KV cache with clean context t_context = torch.ones( [b], device=latents.device, diff --git a/fastvideo/pipelines/basic/wan/stages/dmd.py b/fastvideo/pipelines/basic/wan/stages/dmd.py index 7d8a81e2b..012a89780 100644 --- a/fastvideo/pipelines/basic/wan/stages/dmd.py +++ b/fastvideo/pipelines/basic/wan/stages/dmd.py @@ -42,22 +42,16 @@ class DmdDenoisingStage(DenoisingStage): Returns: The batch with denoised latents. """ - # Setup precision and autocast settings - # TODO(will): make the precision configurable for inference - # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision] target_dtype = torch.bfloat16 autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast - # Get timesteps and calculate warmup steps timesteps = batch.timesteps - # TODO(will): remove this once we add input/output validation for stages if timesteps is None: raise ValueError("Timesteps must be provided") num_inference_steps = batch.num_inference_steps num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - # Prepare image latents and embeddings for I2V generation image_embeds = batch.image_embeds if len(image_embeds) > 0: assert torch.isnan(image_embeds[0]).sum() == 0 @@ -79,7 +73,6 @@ class DmdDenoisingStage(DenoisingStage): }, ) - # Get latents and embeddings assert batch.latents is not None, "latents must be provided" latents = batch.latents @@ -90,10 +83,8 @@ class DmdDenoisingStage(DenoisingStage): dtype=torch.long, device=get_local_torch_device()) - # Run denoising loop with profiler_region("inference_denoising"), self.progress_bar(total=len(timesteps)) as progress_bar: for i, t in enumerate(timesteps): - # Skip if interrupted if hasattr(self, 'interrupt') and self.interrupt: continue # Expand latents for I2V @@ -105,7 +96,6 @@ class DmdDenoisingStage(DenoisingStage): [latent_model_input, batch.image_latent.permute(0, 2, 1, 3, 4)], dim=2).to(target_dtype) assert not torch.isnan(latent_model_input).any(), "latent_model_input contains nan" - # Prepare inputs for transformer t_expand = t.repeat(latent_model_input.shape[0]) guidance_expand = (torch.tensor( [fastvideo_args.pipeline_config.embedded_cfg_scale] * latent_model_input.shape[0], @@ -120,7 +110,6 @@ class DmdDenoisingStage(DenoisingStage): if self.attn_metadata_builder_cls is not None: self.attn_metadata_builder = self.attn_metadata_builder_cls() - # TODO(will): clean this up attn_metadata = self.attn_metadata_builder.build( # type: ignore current_timestep=i, # type: ignore raw_latent_shape=batch.raw_latent_shape[2:5], # type: ignore @@ -140,9 +129,7 @@ class DmdDenoisingStage(DenoisingStage): current_timestep=i, attn_metadata=attn_metadata, forward_batch=batch, - # fastvideo_args=fastvideo_args ): - # Run transformer pred_noise = self.transformer( latent_model_input.permute(0, 2, 1, 3, 4), prompt_embeds, @@ -167,14 +154,11 @@ class DmdDenoisingStage(DenoisingStage): else: latents = pred_video - # Update progress bar if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 and progress_bar is not None): progress_bar.update() - # Gather results if using sequence parallelism latents = latents.permute(0, 2, 1, 3, 4) - # Update batch with final latents batch.latents = latents return batch diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index 24d579adf..765a4788a 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -85,7 +85,6 @@ class ComposedPipelineBase(ABC): training_args: TrainingArgs | None = None fastvideo_args: FastVideoArgs | TrainingArgs | None = None modules: dict[str, Any] = {} - # do not need to include moe related transformers trainable_transformer_names: list[str] = ["transformer"] trainable_transformer_modules: dict[str, torch.nn.Module] = {} post_init_called: bool = False @@ -128,7 +127,6 @@ class ComposedPipelineBase(ABC): *(f"{component_dir}/**" for component_dir in component_dirs), ] - # TODO(will): args should support both inference args and training args def __init__(self, model_path: str, fastvideo_args: FastVideoArgs | TrainingArgs, @@ -168,7 +166,6 @@ class ComposedPipelineBase(ABC): self.local_rank = get_world_group().local_rank - # Load modules directly in initialization logger.info("Loading pipeline modules...") with self.profiler_controller.region("profiler_region_model_loading"): self.modules = self.load_modules(fastvideo_args, loaded_modules) @@ -438,7 +435,6 @@ class ComposedPipelineBase(ABC): else: assert args is not None, "args must be provided for training mode" fastvideo_args = TrainingArgs.from_cli_args(args) - # TODO(will): fix this so that its not so ugly fastvideo_args.model_path = model_path for key, value in kwargs.items(): setattr(fastvideo_args, key, value) @@ -481,7 +477,6 @@ class ComposedPipelineBase(ABC): allow_patterns=self.get_hf_download_allow_patterns(), ) self.model_path = model_path - # fastvideo_args.downloaded_model_path = model_path logger.info("Model path: %s", model_path) config = verify_model_config_and_directory( model_path, @@ -555,7 +550,6 @@ class ComposedPipelineBase(ABC): model_index = self._load_config(self.model_path) logger.info("Loading pipeline modules from config: %s", model_index) - # remove keys that are not pipeline modules model_index.pop("_class_name") model_index.pop("_diffusers_version") model_index.pop("_name_or_path", None) @@ -658,7 +652,6 @@ class ComposedPipelineBase(ABC): logger.warning("Overwriting module %s", module_name) modules[module_name] = module - # Check if all required modules were loaded for module_name in required_modules: if module_name not in modules or modules[module_name] is None: raise ValueError( @@ -764,7 +757,6 @@ class ComposedPipelineBase(ABC): stage_name) self._install_lazy_release_hooks() - # TODO(will): don't hardcode no_grad @torch.no_grad() def forward( self, @@ -785,7 +777,6 @@ class ComposedPipelineBase(ABC): # Execute each stage logger.info("Running pipeline stages: %s", self._stage_name_mapping.keys()) - # logger.info("Batch: %s", batch) try: for stage in self.stages: batch = stage(batch, fastvideo_args) @@ -798,7 +789,6 @@ class ComposedPipelineBase(ABC): self._release_all_lazy_modules() raise - # Return the output return batch def train(self) -> None: diff --git a/fastvideo/pipelines/lora_pipeline.py b/fastvideo/pipelines/lora_pipeline.py index e482ae651..1d1b27f7d 100644 --- a/fastvideo/pipelines/lora_pipeline.py +++ b/fastvideo/pipelines/lora_pipeline.py @@ -230,7 +230,6 @@ class LoRAPipeline(ComposedPipelineBase): # Sequential load (the 32 GB and other cards that cannot hold the text encoder and the DiT at once) omits # the transformer from the first load. The adapter is applied when that module arrives. _constructor_adapter_pending: bool = False - # Set once inference LoRA state exists, so ``add_module`` wraps a transformer that arrives later. _tracks_late_transformers: bool = False # Adapters loaded while the transformer was not in memory; their coverage is checked when it loads. _coverage_deferred: set[str] @@ -250,11 +249,9 @@ class LoRAPipeline(ComposedPipelineBase): self.cur_adapter_path = "" self.cur_adapter_strength = 1.0 self._coverage_deferred = set() - # build list of trainable transformers for transformer_name in self.trainable_transformer_names: if (transformer_name in self.modules and self.modules[transformer_name] is not None): self.trainable_transformer_modules[transformer_name] = (self.modules[transformer_name]) - # check for transformer_2 in case of Wan2.2 MoE or fake_score_transformer_2 if transformer_name.endswith("_2"): raise ValueError( f"trainable_transformer_name override in pipelines should not include _2 suffix: {transformer_name}" @@ -631,7 +628,6 @@ class LoRAPipeline(ComposedPipelineBase): if not self._setting_constructor_adapter and _has_nvfp4_weights_without_bf16( self, self.trainable_transformer_modules): - # TODO(David): Restore the BF16 weights and requantize them after an NVFP4 LoRA adapter change. raise RuntimeError( "Runtime LoRA adapter changes are unsupported after NVFP4 quantization removed the BF16 weights. " "Create a new VideoGenerator with the desired LoRA adapter.") @@ -663,10 +659,8 @@ class LoRAPipeline(ComposedPipelineBase): param_names_mapping_fn = get_param_names_mapping(param_names_mapping) lora_param_names_mapping_fn = get_param_names_mapping(lora_param_names_mapping) - # Extract alpha values and weights in a single pass to_merge_params: defaultdict[Hashable, dict[Any, Any]] = (defaultdict(dict)) for name, weight in lora_state_dict.items(): - # Extract weights (lora_A, lora_B, and lora_alpha) normalized = normalize_lora_key(name) if normalized is None: continue @@ -679,7 +673,6 @@ class LoRAPipeline(ComposedPipelineBase): layer_name = name.replace(".lora_alpha", "") layer_name, _, _ = lora_param_names_mapping_fn(layer_name) target_name, _, _ = param_names_mapping_fn(layer_name) - # Store alpha alongside weights with same target_name base alpha_key = target_name + ".lora_alpha" self.lora_adapters[lora_nickname][alpha_key] = (weight.item() if weight.numel() == 1 else float( weight.mean())) @@ -742,7 +735,6 @@ class LoRAPipeline(ComposedPipelineBase): lora_alpha_name = name + ".lora_alpha" if (lora_A_name in self.lora_adapters[lora_nickname] and lora_B_name in self.lora_adapters[lora_nickname]): - # Get alpha value for this layer (defaults to None if not present) lora_A = self.lora_adapters[lora_nickname][lora_A_name] lora_B = self.lora_adapters[lora_nickname][lora_B_name] # Simple lookup - alpha stored with same naming scheme as lora_A/lora_B @@ -812,7 +804,6 @@ class LoRAPipeline(ComposedPipelineBase): BF16 weight cannot be unmerged and remerged. """ if _has_nvfp4_weights_without_bf16(self, self.trainable_transformer_modules): - # TODO(David): Merge from the packed FP4 weight instead of rejecting. raise RuntimeError( "LoRA merge is unsupported after NVFP4 weight quantization removed the BF16 weights because the " "FP4 weights would not reflect the merged LoRA adapter.") @@ -862,7 +853,6 @@ class LoRAPipeline(ComposedPipelineBase): NVFP4 is rejected: its FP4 buffers cannot be rebuilt from a purged BF16 weight. """ if _has_quantized_nvfp4_weights(self, self.trainable_transformer_modules): - # TODO(David): Preserve BF16 weights and requantize NVFP4 weights after LoRA unmerge. raise RuntimeError( "LoRA unmerge is unsupported after NVFP4 weight quantization because the quantized weights would " "not reflect the unmerged LoRA adapter.") diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index e780d9640..a3244f605 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -67,9 +67,6 @@ class ForwardBatch: execution, allowing methods to update specific components without needing to manage numerous individual parameters. """ - # TODO(will): double check that args are separate from fastvideo_args - # properly. Also maybe think about providing an abstraction for pipeline - # specific arguments. data_type: str generator: torch.Generator | list[torch.Generator] | None = None diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py index 8980d7b7a..e9166dfd3 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/fastvideo/pipelines/pipeline_registry.py @@ -51,7 +51,6 @@ class PipelineType(str, Enum): @dataclass class _PipelineRegistry: # Keyed by pipeline_type -> pipeline_name - # pipelines[pipeline_type][pipeline_name] = pipeline_cls pipelines: dict[str, dict[str, type[ComposedPipelineBase] | None]] = field(default_factory=dict) def get_supported_pipelines(self, pipeline_type: PipelineType) -> Set[str]: diff --git a/fastvideo/pipelines/preprocess/matrixgame2/matrixgame2_preprocess_pipeline.py b/fastvideo/pipelines/preprocess/matrixgame2/matrixgame2_preprocess_pipeline.py index a42af54d6..4ade79923 100644 --- a/fastvideo/pipelines/preprocess/matrixgame2/matrixgame2_preprocess_pipeline.py +++ b/fastvideo/pipelines/preprocess/matrixgame2/matrixgame2_preprocess_pipeline.py @@ -31,7 +31,6 @@ class PreprocessPipeline_MatrixGame2(BasePreprocessPipeline): def get_extra_features(self, valid_data: dict[str, Any], fastvideo_args: FastVideoArgs) -> dict[str, Any]: - # TODO(will): move these to cpu at some point self.get_module("image_encoder").to(get_local_torch_device()) self.get_module("vae").to(get_local_torch_device()) @@ -39,9 +38,6 @@ class PreprocessPipeline_MatrixGame2(BasePreprocessPipeline): """Get CLIP features from the first frame of each video.""" first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute(0, 2, 3, 1) # (B, C, T, H, W) -> (B, H, W, C) _, _, num_frames, height, width = valid_data["pixel_values"].shape - # latent_height = height // self.get_module( - # "vae").spatial_compression_ratio - # latent_width = width // self.get_module("vae").spatial_compression_ratio processed_images = [] # Frame has values between -1 and 1 @@ -51,7 +47,6 @@ class PreprocessPipeline_MatrixGame2(BasePreprocessPipeline): processed_img = self.get_module("image_processor")(images=frame_pil, return_tensors="pt") processed_images.append(processed_img) - # Get CLIP features pixel_values = torch.cat([img['pixel_values'] for img in processed_images], dim=0).to(get_local_torch_device()) with torch.no_grad(): image_inputs = {'pixel_values': pixel_values} @@ -60,7 +55,6 @@ class PreprocessPipeline_MatrixGame2(BasePreprocessPipeline): clip_features = clip_features.last_hidden_state features["clip_feature"] = clip_features - # Get CLIP features from the first frame of each video. video_conditions = [] for frame in first_frame: processed_img = frame.to(device="cpu", dtype=torch.float32) @@ -110,7 +104,6 @@ class PreprocessPipeline_MatrixGame2(BasePreprocessPipeline): # mask_lat_size = mask_lat_size.transpose(1, 2) # mask_lat_size = mask_lat_size.to(latent_condition.device) - # image_latent = torch.concat([mask_lat_size, latent_condition], dim=1) features["first_frame_latent"] = latent_condition diff --git a/fastvideo/pipelines/preprocess/matrixgame2/matrixgame2_preprocess_pipeline_ode_trajectory.py b/fastvideo/pipelines/preprocess/matrixgame2/matrixgame2_preprocess_pipeline_ode_trajectory.py index eb087c01f..4bcd6935d 100644 --- a/fastvideo/pipelines/preprocess/matrixgame2/matrixgame2_preprocess_pipeline_ode_trajectory.py +++ b/fastvideo/pipelines/preprocess/matrixgame2/matrixgame2_preprocess_pipeline_ode_trajectory.py @@ -84,7 +84,6 @@ class PreprocessPipeline_MatrixGame2_ODE_Trajectory(BasePreprocessPipeline): def get_extra_features(self, valid_data: dict[str, Any], fastvideo_args: FastVideoArgs) -> dict[str, Any]: - # TODO(will): move these to cpu at some point self.get_module("image_encoder").to(get_local_torch_device()) self.get_module("vae").to(get_local_torch_device()) @@ -92,9 +91,6 @@ class PreprocessPipeline_MatrixGame2_ODE_Trajectory(BasePreprocessPipeline): """Get CLIP features from the first frame of each video.""" first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute(0, 2, 3, 1) # (B, C, T, H, W) -> (B, H, W, C) _, _, num_frames, height, width = valid_data["pixel_values"].shape - # latent_height = height // self.get_module( - # "vae").spatial_compression_ratio - # latent_width = width // self.get_module("vae").spatial_compression_ratio processed_images = [] # Frame has values between -1 and 1 @@ -104,7 +100,6 @@ class PreprocessPipeline_MatrixGame2_ODE_Trajectory(BasePreprocessPipeline): processed_img = self.get_module("image_processor")(images=frame_pil, return_tensors="pt") processed_images.append(processed_img) - # Get CLIP features pixel_values = torch.cat([img['pixel_values'] for img in processed_images], dim=0).to(get_local_torch_device()) with torch.no_grad(): image_inputs = {'pixel_values': pixel_values} @@ -114,7 +109,6 @@ class PreprocessPipeline_MatrixGame2_ODE_Trajectory(BasePreprocessPipeline): features["clip_feature"] = clip_features features["pil_image"] = first_frame - # Get CLIP features from the first frame of each video. video_conditions = [] for frame in first_frame: processed_img = frame.to(device="cpu", dtype=torch.float32) @@ -224,7 +218,6 @@ class PreprocessPipeline_MatrixGame2_ODE_Trajectory(BasePreprocessPipeline): pixel_values = pixel_values.repeat(1, 1, args.num_frames, 1, 1) valid_data["pixel_values"] = pixel_values - # Get extra features if needed extra_features = self.get_extra_features(valid_data, fastvideo_args) clip_features = extra_features['clip_feature'] @@ -274,7 +267,6 @@ class PreprocessPipeline_MatrixGame2_ODE_Trajectory(BasePreprocessPipeline): trajectory_timesteps.append(result_batch.trajectory_timesteps.cpu()) trajectory_decoded.append(result_batch.trajectory_decoded) - # Prepare extra features extra_features = { "trajectory_latents": trajectory_latents, "trajectory_timesteps": trajectory_timesteps @@ -287,10 +279,8 @@ class PreprocessPipeline_MatrixGame2_ODE_Trajectory(BasePreprocessPipeline): f"decoded_videos/trajectory_decoded_{i}_{j}.mp4", args.train_fps) - # Prepare batch data for Parquet dataset batch_data: list[dict[str, Any]] = [] - # Add progress bar for saving outputs save_pbar = tqdm(enumerate(valid_data["path"]), desc="Saving outputs", unit="item", leave=False) for idx, video_path in save_pbar: @@ -302,7 +292,6 @@ class PreprocessPipeline_MatrixGame2_ODE_Trajectory(BasePreprocessPipeline): keyboard_cond_np = keyboard_cond[idx] if keyboard_cond is not None else None mouse_cond_np = mouse_cond[idx] if mouse_cond is not None else None - # Get trajectory features for this sample traj_latents = extra_features["trajectory_latents"][idx] traj_timesteps = extra_features["trajectory_timesteps"][idx] if isinstance(traj_latents, torch.Tensor): @@ -310,7 +299,6 @@ class PreprocessPipeline_MatrixGame2_ODE_Trajectory(BasePreprocessPipeline): if isinstance(traj_timesteps, torch.Tensor): traj_timesteps = traj_timesteps.cpu().float().numpy() - # Create record for Parquet dataset record: dict[str, Any] = matrixgame2_ode_record_creator(video_name=video_name, clip_feature=clip_feature_np, first_frame_latent=first_frame_latent_np, @@ -354,7 +342,6 @@ class PreprocessPipeline_MatrixGame2_ODE_Trajectory(BasePreprocessPipeline): self.local_rank = int(os.getenv("RANK", 0)) os.makedirs(args.output_dir, exist_ok=True) - # Create directory for combined data self.combined_parquet_dir = os.path.join(args.output_dir, "combined_parquet_dataset") os.makedirs(self.combined_parquet_dir, exist_ok=True) @@ -370,13 +357,11 @@ class PreprocessPipeline_MatrixGame2_ODE_Trajectory(BasePreprocessPipeline): self.preprocess_loader_iter = iter(self.preprocess_dataloader) self.num_processed_samples = 0 - # Add progress bar for video preprocessing self.pbar = tqdm(self.preprocess_loader_iter, desc="Processing videos", unit="batch", disable=self.local_rank != 0) - # Initialize class variables for data sharing self.video_data: dict[str, Any] = {} # Store video metadata and paths self.latent_data: dict[str, Any] = {} # Store latent tensors self.preprocess_action_and_trajectory(fastvideo_args, args) diff --git a/fastvideo/pipelines/preprocess/preprocess_cosmos25_overfit.py b/fastvideo/pipelines/preprocess/preprocess_cosmos25_overfit.py index b1ce99d15..999bd74c3 100644 --- a/fastvideo/pipelines/preprocess/preprocess_cosmos25_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_cosmos25_overfit.py @@ -65,7 +65,6 @@ def main() -> None: os.makedirs(OUTPUT_DIR, exist_ok=True) - # Load captions with open(os.path.join(DATA_DIR, "videos2caption.json")) as f: caption_data = json.load(f) @@ -161,11 +160,9 @@ def main() -> None: } records.append(record) - # Clean up del text_encoder, vae torch.cuda.empty_cache() - # Write parquet table = pa.table( {k: [r[k] for r in records] for k in records[0]}, @@ -175,7 +172,6 @@ def main() -> None: pq.write_table(table, output_path) print(f"\nWrote {len(records)} records to {output_path}") - # Write T2W validation prompts (no image_path for T2W) val_prompts = { "data": [{ "caption": item["cap"][0], diff --git a/fastvideo/pipelines/preprocess/preprocess_cosmos_overfit.py b/fastvideo/pipelines/preprocess/preprocess_cosmos_overfit.py index 24114b8fc..1980cb8d7 100644 --- a/fastvideo/pipelines/preprocess/preprocess_cosmos_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_cosmos_overfit.py @@ -48,7 +48,6 @@ def load_video(path: str, num_frames: int) -> torch.Tensor: cap.release() if len(frames) < num_frames: - # Repeat last frame to fill while len(frames) < num_frames: frames.append(frames[-1]) @@ -66,7 +65,6 @@ def main() -> None: os.makedirs(OUTPUT_DIR, exist_ok=True) - # Load captions with open(os.path.join(DATA_DIR, "videos2caption.json")) as f: caption_data = json.load(f) @@ -150,10 +148,8 @@ def main() -> None: } records.append(record) - # Clean up encoders del text_encoder, tokenizer, vae - # Write parquet table = pa.table( {k: [r[k] for r in records] for k in records[0]}, @@ -163,7 +159,6 @@ def main() -> None: pq.write_table(table, output_path) print(f"\nWrote {len(records)} records to {output_path}") - # Extract first frame from first video as V2W conditioning image import cv2 first_video = os.path.join(DATA_DIR, "videos", caption_data[0]["path"]) cap = cv2.VideoCapture(first_video) diff --git a/fastvideo/pipelines/preprocess/preprocess_hunyuan_overfit.py b/fastvideo/pipelines/preprocess/preprocess_hunyuan_overfit.py index bb4c2ca15..823a4963a 100644 --- a/fastvideo/pipelines/preprocess/preprocess_hunyuan_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_hunyuan_overfit.py @@ -54,7 +54,6 @@ def load_video(path: str, num_frames: int) -> torch.Tensor: cap.release() if len(frames) < num_frames: - # Repeat last frame to fill while len(frames) < num_frames: frames.append(frames[-1]) @@ -73,7 +72,6 @@ def main() -> None: os.makedirs(OUTPUT_DIR, exist_ok=True) - # Load captions with open(os.path.join(DATA_DIR, "videos2caption.json")) as f: caption_data = json.load(f) @@ -177,10 +175,8 @@ def main() -> None: } records.append(record) - # Clean up encoders del llama_enc, llama_tok, clip_enc, clip_tok, vae - # Write parquet table = pa.table( {k: [r[k] for r in records] for k in records[0]}, diff --git a/fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py b/fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py index 1186b48b2..1b8f3e7fe 100644 --- a/fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py @@ -97,7 +97,6 @@ def load_video(path: str, num_frames: int, height: int, width: int) -> torch.Ten "the file is missing, empty, or not readable by OpenCV.") if len(frames) < num_frames: - # Repeat last frame to fill while len(frames) < num_frames: frames.append(frames[-1]) @@ -281,10 +280,8 @@ def main() -> None: } records.append(record) - # Clean up encoders del qwen_enc, qwen_tok, clip_enc, clip_tok, vae - # Write parquet table = pa.table( {k: [r[k] for r in records] for k in records[0]}, diff --git a/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py b/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py index 4eca9a29c..153071201 100644 --- a/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py @@ -220,11 +220,9 @@ def main() -> None: writer.append_data(frame) print(f" Wrote reference clip to {ref_path}") - # Clean up del text_encoder, tokenizer, vae torch.cuda.empty_cache() - # Write parquet (replicated num_copies times; see the comment above num_copies) replicated = [] for copy_idx in range(max(1, num_copies)): for r in records: @@ -241,7 +239,6 @@ def main() -> None: print(f"\nWrote {len(replicated)} records " f"({len(records)} unique x {max(1, num_copies)} copies) to {output_path}") - # Write validation prompts for the validation callback val_prompts = { "data": [{ "caption": (item["cap"][0] if isinstance(item["cap"], list) else item["cap"]), diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py index 9795faded..9420f858b 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py @@ -43,7 +43,6 @@ class BasePreprocessPipeline(ComposedPipelineBase): if not self.post_init_called: self.post_init() - # Initialize class variables for data sharing self.video_data: dict[str, Any] = {} # Store video metadata and paths self.latent_data: dict[str, Any] = {} # Store latent tensors self.preprocess_video_and_text(fastvideo_args, args) @@ -84,7 +83,6 @@ class BasePreprocessPipeline(ComposedPipelineBase): field_filled = False if field.endswith('_bytes'): - # Handle binary tensor data - convert numpy array or tensor to bytes tensor_name = field.replace('_bytes', '') tensor_data = getattr(preprocess_batch, tensor_name, None) if tensor_data is not None: @@ -105,7 +103,6 @@ class BasePreprocessPipeline(ComposedPipelineBase): record[field] = b'' # Empty bytes for missing data elif field.endswith('_shape'): - # Handle tensor shape info tensor_name = field.replace('_shape', '') tensor_data = getattr(preprocess_batch, tensor_name, None) if tensor_data is not None and hasattr(tensor_data, 'shape'): @@ -115,7 +112,6 @@ class BasePreprocessPipeline(ComposedPipelineBase): record[field] = [] elif field.endswith('_dtype'): - # Handle tensor dtype info tensor_name = field.replace('_dtype', '') tensor_data = getattr(preprocess_batch, tensor_name, None) if tensor_data is not None and hasattr(tensor_data, 'dtype'): @@ -125,7 +121,6 @@ class BasePreprocessPipeline(ComposedPipelineBase): record[field] = 'unknown' elif field in ['width', 'height', 'num_frames']: - # Handle integer metadata fields value = getattr(preprocess_batch, field, None) if value is not None: try: @@ -139,8 +134,6 @@ class BasePreprocessPipeline(ComposedPipelineBase): record[field] = 0 elif field in ['duration_sec', 'fps']: - # Handle float metadata fields - # Map schema field names to batch attribute names attr_name = 'duration' if field == 'duration_sec' else field value = getattr(preprocess_batch, attr_name, None) if value is not None: @@ -163,7 +156,6 @@ class BasePreprocessPipeline(ComposedPipelineBase): elif field == 'file_name': attr_name = 'path' elif field == 'id': - # Generate ID from path if available path_value = getattr(preprocess_batch, 'path', None) if path_value: import os @@ -193,7 +185,6 @@ class BasePreprocessPipeline(ComposedPipelineBase): if not field_filled: unfilled_fields.append(field) - # Handle strict mode if strict and unfilled_fields: raise ValueError(f"Required fields were not filled: {unfilled_fields}") @@ -234,12 +225,10 @@ class BasePreprocessPipeline(ComposedPipelineBase): def preprocess_video_and_text(self, fastvideo_args: FastVideoArgs, args): os.makedirs(args.output_dir, exist_ok=True) - # Create directory for combined data combined_parquet_dir = os.path.join(args.output_dir, "combined_parquet_dataset") os.makedirs(combined_parquet_dir, exist_ok=True) local_rank = int(os.getenv("RANK", 0)) - # Get how many samples have already been processed start_idx = 0 for root, _, files in os.walk(combined_parquet_dir): for file in files: @@ -257,7 +246,6 @@ class BasePreprocessPipeline(ComposedPipelineBase): ) num_processed_samples = 0 - # Add progress bar for video preprocessing pbar = tqdm(train_dataloader, desc="Processing videos", unit="batch", disable=local_rank != 0) for batch_idx, data in enumerate(pbar): @@ -292,7 +280,6 @@ class BasePreprocessPipeline(ComposedPipelineBase): latents = self.get_module("vae").encode(valid_data["pixel_values"].to( get_local_torch_device())).mean - # Get extra features if needed extra_features = self.get_extra_features(valid_data, fastvideo_args) batch_captions = valid_data["text"] @@ -309,13 +296,11 @@ class BasePreprocessPipeline(ComposedPipelineBase): 0], result_batch.prompt_attention_mask[0] assert prompt_embeds.shape[0] == prompt_attention_mask.shape[0] - # Get sequence lengths from attention masks (number of 1s) seq_lens = prompt_attention_mask.sum(dim=1) non_padded_embeds = [] non_padded_masks = [] - # Process each item in the batch for i in range(prompt_embeds.size(0)): seq_len = seq_lens[i].item() # Slice the embeddings and masks to keep only non-padding parts @@ -329,21 +314,16 @@ class BasePreprocessPipeline(ComposedPipelineBase): bs = len(valid_indices) prompt_embeds = [torch.zeros(0) for _ in range(bs)] - # Prepare batch data for Parquet dataset batch_data = [] - # Add progress bar for saving outputs save_pbar = tqdm(enumerate(valid_data["path"]), desc="Saving outputs", unit="item", leave=False) for idx, video_path in save_pbar: - # Get the corresponding latent and info using video name latent = latents[idx].cpu() video_name = os.path.basename(video_path).split(".")[0] - # Convert tensors to numpy arrays vae_latent = latent.cpu().numpy() text_embedding = prompt_embeds[idx].cpu().float().numpy() - # Get extra features for this sample if needed sample_extra_features = {} if extra_features: for key, value in extra_features.items(): @@ -352,7 +332,6 @@ class BasePreprocessPipeline(ComposedPipelineBase): else: sample_extra_features[key] = value[idx] - # Create record for Parquet dataset record = self.create_record(video_name=video_name, vae_latent=vae_latent, text_embedding=text_embedding, diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_i2v.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_i2v.py index 0ee73a9e2..02e55c16e 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_i2v.py +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_i2v.py @@ -44,14 +44,12 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline): def get_extra_features(self, valid_data: dict[str, Any], fastvideo_args: FastVideoArgs) -> dict[str, Any]: - # TODO(will): move these to cpu at some point self.get_module("vae").to(get_local_torch_device()) features = {} first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute(0, 2, 3, 1) # (B, C, T, H, W) -> (B, H, W, C) _, _, num_frames, height, width = valid_data["pixel_values"].shape - # Get CLIP features if image_encoder is available if (self.get_module("image_encoder") is not None and self.get_module("image_processor") is not None): self.get_module("image_encoder").to(get_local_torch_device()) @@ -72,7 +70,6 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline): features["clip_feature"] = clip_features - # Get VAE features from the first frame of each video video_conditions = [] for frame in first_frame: processed_img = frame.to(device="cpu", dtype=torch.float32) diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py index a88ae5539..3c53aea97 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py @@ -100,7 +100,6 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): "path": [data["path"][i] for i in valid_indices], } - # Add fps and duration if available in data if "fps" in data: valid_data["fps"] = [data["fps"][i] for i in valid_indices] if "duration" in data: @@ -183,19 +182,15 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): f"decoded_videos/trajectory_decoded_{i}_{j}.mp4", args.train_fps) - # Prepare batch data for Parquet dataset batch_data: list[dict[str, Any]] = [] - # Add progress bar for saving outputs save_pbar = tqdm(enumerate(valid_data["path"]), desc="Saving outputs", unit="item", leave=False) for idx, video_path in save_pbar: video_name = os.path.basename(video_path).split(".")[0] - # Convert tensors to numpy arrays text_embedding = prompt_embeds[idx].cpu().numpy() - # Get extra features for this sample sample_extra_features = {} if extra_features: for key, value in extra_features.items(): @@ -250,7 +245,6 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): self.local_rank = int(os.getenv("RANK", 0)) os.makedirs(args.output_dir, exist_ok=True) - # Create directory for combined data self.combined_parquet_dir = os.path.join(args.output_dir, "combined_parquet_dataset") os.makedirs(self.combined_parquet_dir, exist_ok=True) @@ -266,13 +260,11 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): self.preprocess_loader_iter = iter(self.preprocess_dataloader) self.num_processed_samples = 0 - # Add progress bar for video preprocessing self.pbar = tqdm(self.preprocess_loader_iter, desc="Processing videos", unit="batch", disable=self.local_rank != 0) - # Initialize class variables for data sharing self.video_data: dict[str, Any] = {} # Store video metadata and paths self.latent_data: dict[str, Any] = {} # Store latent tensors self.preprocess_text_and_trajectory(fastvideo_args, args) diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_text.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_text.py index 241ecc5b1..edcf70e10 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_text.py +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_text.py @@ -89,16 +89,13 @@ class PreprocessPipeline_Text(BasePreprocessPipeline): logger.info("===== prompt_embeds: %s", prompt_embeds.shape) logger.info("===== prompt_attention_masks: %s", prompt_attention_masks.shape) - # Prepare batch data for Parquet dataset batch_data = [] - # Add progress bar for saving outputs save_pbar = tqdm(enumerate(valid_data["path"]), desc="Saving outputs", unit="item", leave=False) for idx, text_path in save_pbar: text_name = os.path.basename(text_path).split(".")[0] - # Convert tensors to numpy arrays text_embedding = prompt_embeds[idx].cpu().numpy() # Create record for Parquet dataset (text-only schema) @@ -143,7 +140,6 @@ class PreprocessPipeline_Text(BasePreprocessPipeline): self.local_rank = int(os.getenv("RANK", 0)) os.makedirs(args.output_dir, exist_ok=True) - # Create directory for combined data self.combined_parquet_dir = os.path.join(args.output_dir, "combined_parquet_dataset") os.makedirs(self.combined_parquet_dir, exist_ok=True) @@ -159,13 +155,11 @@ class PreprocessPipeline_Text(BasePreprocessPipeline): self.preprocess_loader_iter = iter(self.preprocess_dataloader) self.num_processed_samples = 0 - # Add progress bar for text preprocessing self.pbar = tqdm(self.preprocess_loader_iter, desc="Processing text", unit="batch", disable=self.local_rank != 0) - # Initialize class variables for data sharing self.text_data: dict[str, Any] = {} # Store text metadata and paths self.preprocess_text_only(fastvideo_args, args) diff --git a/fastvideo/pipelines/stages/base.py b/fastvideo/pipelines/stages/base.py index 71d9dd013..ab32588b0 100644 --- a/fastvideo/pipelines/stages/base.py +++ b/fastvideo/pipelines/stages/base.py @@ -95,7 +95,6 @@ class PipelineStage(ABC): if not verification_result.is_valid(): failed_fields = verification_result.get_failed_fields() if failed_fields: - # Get detailed failure information detailed_summary = verification_result.get_failure_summary() failed_fields_str = ", ".join(failed_fields) @@ -138,7 +137,6 @@ class PipelineStage(ABC): stage_key = getattr(self, "_pipeline_stage_name", stage_class_name) stage_name = f"{stage_key}|{stage_class_name}" - # Check if verification is enabled (simple approach for prototype) enable_verification = getattr(fastvideo_args, 'enable_stage_verification', False) if enable_verification: diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/pipelines/stages/decoding.py index 27b5160a3..1928d433f 100644 --- a/fastvideo/pipelines/stages/decoding.py +++ b/fastvideo/pipelines/stages/decoding.py @@ -180,8 +180,6 @@ class DecodingStage(PipelineStage): with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): if fastvideo_args.pipeline_config.vae_tiling: self.vae.enable_tiling() - # if fastvideo_args.vae_sp: - # self.vae.enable_parallel() if not vae_autocast_enabled: latents = latents.to(vae_dtype) # Flux2's image VAE expects 4D (B, C, H, W); squeeze the singleton T @@ -206,7 +204,6 @@ class DecodingStage(PipelineStage): if squeezed_for_vae: image = image.unsqueeze(2) - # Normalize image to [0, 1] range image = (image / 2 + 0.5).clamp(0, 1) return image @@ -241,7 +238,6 @@ class DecodingStage(PipelineStage): latents = self._denormalize_latents(latents, fastvideo_args) - # Initialize cache if needed if cache is None: cache = self.vae.get_streaming_cache() @@ -253,7 +249,6 @@ class DecodingStage(PipelineStage): latents = latents.to(vae_dtype) image, cache = self.vae.streaming_decode(latents, cache, is_first_chunk) - # Normalize image to [0, 1] range image = (image / 2 + 0.5).clamp(0, 1) assert cache is not None, "cache should not be None after streaming_decode" return image, cache @@ -288,7 +283,6 @@ class DecodingStage(PipelineStage): - output: Decoded frames (batch, channels, frames, height, width) as CPU float32 - trajectory_decoded (if requested): List of decoded frames per timestep """ - # load vae if not already loaded (used for memory constrained devices) pipeline = self.pipeline() if self.pipeline else None if not fastvideo_args.model_loaded["vae"]: loader = VAELoader() @@ -317,10 +311,8 @@ class DecodingStage(PipelineStage): decoded_frames = self.decode(cur_latent, fastvideo_args) batch.trajectory_decoded.append(decoded_frames.cpu().float()) - # Convert to float32 for compatibility frames = frames.to(torch.float32) - # Update batch with decoded image batch.output = frames # Offload models if needed diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 4e8faa0ee..04ca2ea09 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -105,7 +105,6 @@ class DenoisingStage(PipelineStage): pipeline.add_module("transformer", self.transformer) fastvideo_args.model_loaded["transformer"] = True - # Prepare extra step kwargs for scheduler extra_step_kwargs = self.prepare_extra_func_kwargs( self.scheduler.step, { @@ -114,9 +113,6 @@ class DenoisingStage(PipelineStage): }, ) - # Setup precision and autocast settings - # TODO(will): make the precision configurable for inference - # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision] target_dtype = torch.bfloat16 # Flux2-only denoising compensations. # @@ -149,15 +145,12 @@ class DenoisingStage(PipelineStage): scheduler_fp32 = getattr(fastvideo_args.pipeline_config, "scheduler_step_in_fp32", False) local_device = get_local_torch_device() - # Get timesteps and calculate warmup steps timesteps = batch.timesteps - # TODO(will): remove this once we add input/output validation for stages if timesteps is None: raise ValueError("Timesteps must be provided") num_inference_steps = batch.num_inference_steps num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order - # Prepare image latents and embeddings for I2V generation image_embeds = batch.image_embeds if len(image_embeds) > 0: assert not torch.isnan(image_embeds[0]).any(), "image_embeds contains nan" @@ -211,7 +204,6 @@ class DenoisingStage(PipelineStage): animate_kwargs, animate_uncond_kwargs = self._animate_conditioning_kwargs(batch) - # Get latents and embeddings latents = batch.latents cast_embeds = getattr(fastvideo_args.pipeline_config.dit_config, "cast_prompt_embeds_to_dit_dtype", False) if cast_embeds: @@ -238,7 +230,6 @@ class DenoisingStage(PipelineStage): latents = state.latents model_kwargs, model_kwargs_uncond = self.prepare_model_kwargs(batch, fastvideo_args) - # Initialize lists for ODE trajectory trajectory_timesteps: list[torch.Tensor] = [] trajectory_latents: list[torch.Tensor] = [] @@ -301,10 +292,8 @@ class DenoisingStage(PipelineStage): _cfg_gate_reused_delta = 0 _cfg_gate_invalidations = 0 - # Run denoising loop with profiler_region("inference_denoising"), self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): - # Skip if interrupted if hasattr(self, 'interrupt') and self.interrupt: continue @@ -358,7 +347,6 @@ class DenoisingStage(PipelineStage): if self.attn_metadata_builder_cls is not None: self.attn_metadata_builder = self.attn_metadata_builder_cls() - # TODO(will): clean this up attn_metadata = self.attn_metadata_builder.build( # type: ignore current_timestep=i, # type: ignore raw_latent_shape=batch.raw_latent_shape[2:5], # type: ignore @@ -374,7 +362,6 @@ class DenoisingStage(PipelineStage): self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls() if self.attn_metadata_builder_cls is not None: self.attn_metadata_builder = self.attn_metadata_builder_cls() - # Prepare V-MoBA parameters from config moba_params = fastvideo_args.moba_config.copy() moba_params.update({ "current_timestep": i, @@ -388,18 +375,12 @@ class DenoisingStage(PipelineStage): attn_metadata = None else: attn_metadata = None - # TODO(will): finalize the interface. vLLM uses this to - # support torch dynamo compilation. They pass in - # attn_metadata, vllm_config, and num_tokens. We can pass in - # fastvideo_args or training_args, and attn_metadata. batch.is_cfg_negative = False with set_forward_context( current_timestep=i, attn_metadata=attn_metadata, forward_batch=batch, - # fastvideo_args=fastvideo_args ): - # Run transformer noise_pred = current_model( latent_model_input, prompt_embeds, @@ -471,7 +452,6 @@ class DenoisingStage(PipelineStage): noise_pred = noise_pred_uncond + current_guidance_scale * (noise_pred_text - noise_pred_uncond) - # Apply guidance rescale if needed if batch.guidance_rescale > 0.0: # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf noise_pred = self.rescale_noise_cfg( @@ -487,12 +467,10 @@ class DenoisingStage(PipelineStage): latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] latents = self.finish_step(latents, state) - # save trajectory latents if needed if batch.return_trajectory_latents: trajectory_timesteps.append(t) trajectory_latents.append(latents) - # Update progress bar if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 and progress_bar is not None): progress_bar.update() @@ -525,7 +503,6 @@ class DenoisingStage(PipelineStage): batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu() batch.trajectory_latents = trajectory_tensor.cpu() - # Update batch with final latents batch.latents = latents if fastvideo_args.dit_layerwise_offload: @@ -911,8 +888,6 @@ class CosmosDenoisingStage(DenoisingStage): else: final_x0 = cond_x0 - # Convert x0 to velocity for - # FlowMatchEulerDiscreteScheduler. velocity = (latents - final_x0) / sigma.clamp(min=1e-6) latents = self.scheduler.step( @@ -1055,7 +1030,6 @@ class Cosmos25DenoisingStage(CosmosDenoisingStage): with self.progress_bar(total=len(timesteps)) as progress_bar: for i, t in enumerate(timesteps): - # Stop if interrupted if getattr(self, "interrupt", False): break t_val = float(t) diff --git a/fastvideo/pipelines/stages/encoding.py b/fastvideo/pipelines/stages/encoding.py index 13f2c1feb..4901b7b54 100644 --- a/fastvideo/pipelines/stages/encoding.py +++ b/fastvideo/pipelines/stages/encoding.py @@ -67,23 +67,18 @@ class EncodingStage(PipelineStage): vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast - # Normalize input to [-1, 1] range (reverse of decoding normalization) latents = (batch.latents * 2.0 - 1.0).clamp(-1, 1) - # Move to appropriate device and dtype latents = latents.to(get_local_torch_device()) # Encode image to latents with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): if fastvideo_args.pipeline_config.vae_tiling: self.vae.enable_tiling() - # if fastvideo_args.vae_sp: - # self.vae.enable_parallel() if not vae_autocast_enabled: latents = latents.to(vae_dtype) latents = self.vae.encode(latents).mean - # Update batch with encoded latents batch.latents = latents # Offload models if needed diff --git a/fastvideo/pipelines/stages/flux_stages.py b/fastvideo/pipelines/stages/flux_stages.py index d38502c56..b5e465182 100644 --- a/fastvideo/pipelines/stages/flux_stages.py +++ b/fastvideo/pipelines/stages/flux_stages.py @@ -292,7 +292,6 @@ class FluxDenoisingStage(PipelineStage): step_extras = self._step_kwargs(self.scheduler.step, batch) for t in timesteps: - # Stop if interrupted if getattr(self, "interrupt", False): break t_scalar = t diff --git a/fastvideo/pipelines/stages/hyworld_denoising.py b/fastvideo/pipelines/stages/hyworld_denoising.py index 763309da8..0a192dbbd 100644 --- a/fastvideo/pipelines/stages/hyworld_denoising.py +++ b/fastvideo/pipelines/stages/hyworld_denoising.py @@ -74,7 +74,6 @@ class HYWorldDenoisingStage(DenoisingStage): pipeline.add_module("transformer", self.transformer) fastvideo_args.model_loaded["transformer"] = True - # Extract HYWorld-specific parameters from batch.extra or batch attributes viewmats = getattr(batch, "viewmats", None) or batch.extra.get("viewmats", None) Ks = getattr(batch, "Ks", None) or batch.extra.get("Ks", None) action = getattr(batch, "action", None) or batch.extra.get("action", None) @@ -91,23 +90,19 @@ class HYWorldDenoisingStage(DenoisingStage): "Provide pose in batch.pose or batch.extra['pose'], or " "viewmats/Ks/action in batch.extra.") - # Get num_frames from batch num_frames = batch.num_frames if isinstance(num_frames, list): num_frames = num_frames[0] - # Calculate number of latents and convert pose to tensors latent_num = compute_latent_num(num_frames) viewmats, Ks, action = pose_to_input(pose, latent_num) - # Add batch dimension viewmats = viewmats.unsqueeze(0) # (1, T, 4, 4) Ks = Ks.unsqueeze(0) # (1, T, 3, 3) action = action.unsqueeze(0) # (1, T) logger.info("Converted pose '%s' to viewmats, Ks, and action for %d latent frames", pose, latent_num) - # Prepare extra step kwargs for scheduler extra_step_kwargs = self.prepare_extra_func_kwargs( self.scheduler.step, { @@ -120,14 +115,12 @@ class HYWorldDenoisingStage(DenoisingStage): target_dtype = torch.bfloat16 autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast - # Get timesteps and calculate warmup steps timesteps = batch.timesteps if timesteps is None: raise ValueError("Timesteps must be provided") num_inference_steps = batch.num_inference_steps num_warmup_steps = (len(timesteps) - num_inference_steps * self.scheduler.order) - # Prepare image latents and embeddings for I2V generation image_embeds = batch.image_embeds if len(image_embeds) > 0: assert not torch.isnan(image_embeds[0]).any(), "image_embeds contains nan" @@ -141,7 +134,6 @@ class HYWorldDenoisingStage(DenoisingStage): }, ) - # Get latents and embeddings latents = batch.latents prompt_embeds = batch.prompt_embeds assert not torch.isnan(prompt_embeds[0]).any(), "prompt_embeds contains nan" @@ -154,7 +146,6 @@ class HYWorldDenoisingStage(DenoisingStage): assert latent_model_input.shape[0] == 1, "only support batch size 1" device = get_local_torch_device() - # Generate local points if not provided if points_local is None: points_local = generate_points_in_sphere(50000, 8.0).to(device) else: @@ -164,18 +155,15 @@ class HYWorldDenoisingStage(DenoisingStage): # batch.image_latent is already [1, 33, T, H, W] with first frame encoded, rest zeros cond_latents = batch.image_latent - # Calculate chunk configuration latent_frames = latents.shape[2] chunk_num = latent_frames // chunk_latent_frames - # Initialize lists for ODE trajectory trajectory_timesteps: list[torch.Tensor] = [] trajectory_latents: list[torch.Tensor] = [] # Main chunk processing loop for chunk_i in range(chunk_num): if chunk_i > 0: - # Select context frames based on camera alignment current_frame_idx = chunk_i * chunk_latent_frames selected_frame_indices = [] @@ -196,11 +184,9 @@ class HYWorldDenoisingStage(DenoisingStage): selected_frame_indices.extend(selected_history_frame_id) selected_frame_indices = sorted(list(set(selected_frame_indices))) - # Remove current chunk frames from context to_remove = list(range(current_frame_idx, current_frame_idx + chunk_latent_frames)) selected_frame_indices = [x for x in selected_frame_indices if x not in to_remove] - # Extract context frames context_latents = latents[:, :, selected_frame_indices] context_w2c = viewmats[:, selected_frame_indices] context_Ks = Ks[:, selected_frame_indices] @@ -208,14 +194,12 @@ class HYWorldDenoisingStage(DenoisingStage): self.scheduler.set_timesteps(num_inference_steps, device=device) - # Define chunk boundaries start_idx = chunk_i * chunk_latent_frames end_idx = chunk_i * chunk_latent_frames + chunk_latent_frames # Denoising loop for this chunk with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): - # Stop if interrupted if getattr(self, "interrupt", False): break @@ -249,7 +233,6 @@ class HYWorldDenoisingStage(DenoisingStage): latent_model_input = torch.cat([context_latents, latents_model_now], dim=2) cond_latents_input = cond_latents[:, :, :latent_model_input.shape[2]] - # Prepare viewmats, Ks, action for current chunk viewmats_input = viewmats[:, start_idx:end_idx] Ks_input = Ks[:, start_idx:end_idx] action_input = action[:, start_idx:end_idx] @@ -259,7 +242,6 @@ class HYWorldDenoisingStage(DenoisingStage): Ks_input = torch.cat([context_Ks, Ks_input], dim=1) action_input = torch.cat([context_action, action_input], dim=1) - # Prepare latent input (concatenate with cond_latents if needed) latents_concat = torch.concat([latent_model_input, cond_latents_input], dim=1) latents_concat = self.scheduler.scale_model_input(latents_concat, t) @@ -278,7 +260,6 @@ class HYWorldDenoisingStage(DenoisingStage): current_model = self.transformer batch.is_cfg_negative = False - # Prepare transformer kwargs with HYWorld-specific inputs transformer_kwargs = { **image_kwargs, "timestep": t_expand, @@ -288,7 +269,6 @@ class HYWorldDenoisingStage(DenoisingStage): "action": action_input.to(target_dtype), } - # Set encoder_attention_mask for positive/negative conditioning pos_transformer_kwargs = { **transformer_kwargs, "encoder_attention_mask": batch.prompt_attention_mask } @@ -324,7 +304,6 @@ class HYWorldDenoisingStage(DenoisingStage): noise_pred = noise_pred_uncond + batch.guidance_scale * (noise_pred_text - noise_pred_uncond) - # Apply guidance rescale if needed if batch.guidance_rescale > 0.0: noise_pred = self.rescale_noise_cfg( noise_pred, @@ -342,17 +321,14 @@ class HYWorldDenoisingStage(DenoisingStage): # Update only the current chunk's latents latents[:, :, start_idx:end_idx] = latent_model_input[:, :, -chunk_latent_frames:] - # Save trajectory latents if needed if batch.return_trajectory_latents: trajectory_timesteps.append(t) trajectory_latents.append(latents.clone()) - # Update progress bar if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 and progress_bar is not None): progress_bar.update() - # Handle trajectory output trajectory_tensor: torch.Tensor | None = None if trajectory_latents: trajectory_tensor = torch.stack(trajectory_latents, dim=1) @@ -365,7 +341,6 @@ class HYWorldDenoisingStage(DenoisingStage): batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu() batch.trajectory_latents = trajectory_tensor.cpu() - # Update batch with final latents batch.latents = latents return batch diff --git a/fastvideo/pipelines/stages/image_encoding.py b/fastvideo/pipelines/stages/image_encoding.py index 6ebb6db86..b0061a6af 100644 --- a/fastvideo/pipelines/stages/image_encoding.py +++ b/fastvideo/pipelines/stages/image_encoding.py @@ -266,7 +266,6 @@ class HYWorldImageEncodingStage(ImageEncodingStage): num_vision_tokens = 729 # (384/14)^2 for SigLIP vision_dim = 1152 # SigLIP hidden size - # Get temporal dimension from raw_latent_shape (set by LatentPreparationStage) raw_latent_shape = list(batch.raw_latent_shape) latent_channels = raw_latent_shape[1] latent_temporal = raw_latent_shape[2] # T dimension @@ -331,7 +330,6 @@ class HYWorldImageEncodingStage(ImageEncodingStage): if isinstance(image, np.ndarray): image = PILImage.fromarray(image) - # Get target size from batch origin_size = image.size target_height, target_width = batch.height, batch.width @@ -388,7 +386,6 @@ class HYWorldImageEncodingStage(ImageEncodingStage): latent_width, device=device) - # Initialize video latent placeholder raw_latent_shape[1] = 1 batch.video_latent = torch.zeros(tuple(raw_latent_shape), device=device) @@ -536,7 +533,6 @@ class ImageVAEEncodingStage(PipelineStage): self.vae = self.vae.to(get_local_torch_device()) - # Process single image for I2V latent_height = height // self.vae.spatial_compression_ratio latent_width = width // self.vae.spatial_compression_ratio image = batch.pil_image @@ -560,8 +556,6 @@ class ImageVAEEncodingStage(PipelineStage): with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): if fastvideo_args.pipeline_config.vae_tiling: self.vae.enable_tiling() - # if fastvideo_args.vae_sp: - # self.vae.enable_parallel() if not vae_autocast_enabled: video_condition = video_condition.to(vae_dtype) encoder_output = self.vae.encode(video_condition) @@ -574,7 +568,6 @@ class ImageVAEEncodingStage(PipelineStage): raise ValueError("Generator must be provided") latent_condition = self.retrieve_latents(encoder_output, generator) - # Apply shifting if needed if (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None): if isinstance(self.vae.shift_factor, torch.Tensor): latent_condition -= self.vae.shift_factor.to(latent_condition.device, latent_condition.dtype) @@ -724,7 +717,6 @@ class VideoVAEEncodingStage(ImageVAEEncodingStage): self.vae = self.vae.to(get_local_torch_device()) - # Prepare video tensor from control video video_condition = self._prepare_control_video_tensor(batch.video_latent, num_frames, height, width).to(get_local_torch_device(), dtype=torch.float32) @@ -788,7 +780,6 @@ class VideoVAEEncodingStage(ImageVAEEncodingStage): else: video_tensor = torch.zeros(1, 3, 0, height, width, device=get_local_torch_device(), dtype=torch.float32) elif isinstance(control_video, torch.Tensor): - # Handle tensor input [batch, channels, frames, height, width] video_tensor = control_video.to(get_local_torch_device(), dtype=torch.float32) if video_tensor.shape[2] > num_frames: @@ -873,14 +864,12 @@ class MatrixGame2ImageVAEEncodingStage(ImageVAEEncodingStage): image = batch.pil_image - # Handle tensor input from causal pipeline if isinstance(image, torch.Tensor): # Causal pipeline provides tensor in [B, C, F, H, W] format if image.dim() == 5: # Already 5D, extract first frame for conditioning # Shape: [B, C, F, H, W] -> use first frame [B, C, 1, H, W] first_frame = image[:, :, :1] # Keep dim, [B, C, 1, H, W] - # Create video condition with first frame + zeros video_condition = torch.cat([ first_frame, first_frame.new_zeros(first_frame.shape[0], first_frame.shape[1], num_frames - 1, @@ -908,7 +897,6 @@ class MatrixGame2ImageVAEEncodingStage(ImageVAEEncodingStage): # (B, C, H, W) -> (B, C, 1, H, W) image = image.unsqueeze(2) - # Create video tensor with first frame as image, rest as zeros video_condition = torch.cat([ image, image.new_zeros(image.shape[0], image.shape[1], num_frames - 1, image.shape[3], image.shape[4]) @@ -934,14 +922,12 @@ class MatrixGame2ImageVAEEncodingStage(ImageVAEEncodingStage): # manually using latents_mean and latents_std from config... if (hasattr(self.vae.config, 'latents_mean') and hasattr(self.vae.config, 'latents_std')): - # Convert config values to tensors latents_mean = torch.tensor(self.vae.config.latents_mean, device=img_cond.device, dtype=img_cond.dtype).view(1, -1, 1, 1, 1) latents_std = torch.tensor(self.vae.config.latents_std, device=img_cond.device, dtype=img_cond.dtype).view(1, -1, 1, 1, 1) - # Apply normalization: (latent - mean) * (1/std) img_cond = (img_cond - latents_mean) / latents_std elif (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None): # Fallback to shift_factor/scaling_factor if available diff --git a/fastvideo/pipelines/stages/input_validation.py b/fastvideo/pipelines/stages/input_validation.py index d6fd0a337..2529934c0 100644 --- a/fastvideo/pipelines/stages/input_validation.py +++ b/fastvideo/pipelines/stages/input_validation.py @@ -59,31 +59,24 @@ class InputValidationStage(PipelineStage): self._generate_seeds(batch, fastvideo_args) - # Ensure prompt is properly formatted if batch.prompt is None and batch.prompt_embeds is None: raise ValueError("Either `prompt` or `prompt_embeds` must be provided") - # Ensure negative prompt is properly formatted if using classifier-free guidance if (batch.do_classifier_free_guidance and batch.negative_prompt is None and not batch.negative_prompt_embeds): raise ValueError("For classifier-free guidance, either `negative_prompt` or " "`negative_prompt_embeds` must be provided") - # Validate height and width if batch.height is None or batch.width is None: raise ValueError("Height and width must be provided. Please set `height` and `width`.") if batch.height % 8 != 0 or batch.width % 8 != 0: raise ValueError(f"Height and width must be divisible by 8 but are {batch.height} and {batch.width}.") - # Validate number of inference steps if batch.num_inference_steps <= 0: raise ValueError(f"Number of inference steps must be positive, but got {batch.num_inference_steps}") - # Validate guidance scale if using classifier-free guidance if batch.do_classifier_free_guidance and batch.guidance_scale <= 0: raise ValueError(f"Guidance scale must be positive, but got {batch.guidance_scale}") - # for i2v, get image from image_path - # @TODO(Wei) hard-coded for wan2.2 5b ti2v for now. Should put this in image_encoding stage if batch.image_path is not None: if batch.image_path.endswith(".mp4"): image = load_video(batch.image_path)[0] @@ -132,7 +125,6 @@ class InputValidationStage(PipelineStage): pil_images, original_fps = load_video(batch.video_path, return_fps=True) logger.info("Loaded video with %s frames, original FPS: %s", len(pil_images), original_fps) - # Get target parameters from batch target_fps = batch.fps target_num_frames = batch.num_frames target_height = batch.height @@ -157,7 +149,6 @@ class InputValidationStage(PipelineStage): resized_img = resize(pil_img, target_height, target_width, resize_mode="default", resample="lanczos") resized_images.append(resized_img) - # Convert PIL images to numpy array video_numpy = pil_to_numpy(resized_images) video_numpy = normalize(video_numpy) video_tensor = numpy_to_pt(video_numpy) @@ -167,7 +158,6 @@ class InputValidationStage(PipelineStage): batch.video_latent = input_video - # Validate action control inputs (Matrix-Game) if batch.mouse_cond is not None: if batch.mouse_cond.dim() != 3 or batch.mouse_cond.shape[-1] != 2: raise ValueError(f"mouse_cond must have shape (B, T, 2), but got {batch.mouse_cond.shape}") diff --git a/fastvideo/pipelines/stages/kandinsky6_sr.py b/fastvideo/pipelines/stages/kandinsky6_sr.py index 1a711e2dc..e03c3d9b9 100644 --- a/fastvideo/pipelines/stages/kandinsky6_sr.py +++ b/fastvideo/pipelines/stages/kandinsky6_sr.py @@ -263,7 +263,6 @@ class Kandinsky6SRDenoisingStage(PipelineStage): state, cond = x[..., :channels], x[..., channels:] self._set_timesteps(num_steps, device, use_piflow) for i, t in enumerate(self.scheduler.timesteps): - # Stop if interrupted if getattr(self, "interrupt", False): break timestep = t.to(device=device, dtype=torch.float32).expand(x.shape[0]) diff --git a/fastvideo/pipelines/stages/latent_preparation.py b/fastvideo/pipelines/stages/latent_preparation.py index e0116af9a..035d0d786 100644 --- a/fastvideo/pipelines/stages/latent_preparation.py +++ b/fastvideo/pipelines/stages/latent_preparation.py @@ -74,7 +74,6 @@ class LatentPreparationStage(PipelineStage): # Adjust batch size for number of videos per prompt batch_size *= batch.num_videos_per_prompt - # Get required parameters if not batch.prompt_embeds: # Create a dummy zero-length text embedding to satisfy downstream checks. # Matrix-Game models have text_dim=0 and ignore encoder_hidden_states. @@ -96,11 +95,9 @@ class LatentPreparationStage(PipelineStage): height = batch.height width = batch.width - # TODO(will): remove this once we add input/output validation for stages if height is None or width is None: raise ValueError("Height and width must be provided") - # Calculate latent shape bcthw_shape: tuple[int, ...] | None = None if self.use_btchw_layout: shape = ( @@ -121,12 +118,10 @@ class LatentPreparationStage(PipelineStage): ) bcthw_shape = shape - # Validate generator if it's a list if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" f" size of {batch_size}. Make sure the batch size matches the length of the generators.") - # Generate or use provided latents if latents is None: latents = randn_tensor( shape, @@ -141,7 +136,6 @@ class LatentPreparationStage(PipelineStage): if hasattr(self.scheduler, "init_noise_sigma"): latents = latents * self.scheduler.init_noise_sigma - # Update batch with prepared latents batch.latents = latents batch.raw_latent_shape = bcthw_shape @@ -201,7 +195,6 @@ class CosmosLatentPreparationStage(PipelineStage): # Adjust batch size for number of videos per prompt batch_size *= batch.num_videos_per_prompt - # Get required parameters # Force float32 for latent preparation dtype = torch.float32 device = get_local_torch_device() @@ -240,12 +233,10 @@ class CosmosLatentPreparationStage(PipelineStage): processed_image = image_processor.preprocess(batch.pil_image, height, width) - # Add time dimension video = processed_image.unsqueeze(2) video = video.to(device=device, dtype=torch.bfloat16) elif hasattr(batch, 'preprocessed_image') and batch.preprocessed_image is not None: - # Convert preprocessed image to video format if isinstance(batch.preprocessed_image, torch.Tensor): if batch.preprocessed_image.dim() == 4: # [B, C, H, W] -> [B, C, T, H, W] video = batch.preprocessed_image.unsqueeze(2) @@ -301,7 +292,6 @@ class CosmosLatentPreparationStage(PipelineStage): init_latents = torch.cat(init_latents, dim=0).to(dtype) - # Apply latent normalization if hasattr(self.vae.config, 'latents_mean') and hasattr(self.vae.config, 'latents_std'): latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, self.vae.config.z_dim, 1, 1, 1).to(device, dtype) @@ -316,9 +306,7 @@ class CosmosLatentPreparationStage(PipelineStage): else: num_cond_latent_frames = 0 - # Generate or use provided latents if latents is None: - # Use float32 for randn_tensor latents = randn_tensor(shape, generator=torch.Generator(device="cpu").manual_seed(200), device=device, diff --git a/fastvideo/pipelines/stages/matrixgame2_denoising.py b/fastvideo/pipelines/stages/matrixgame2_denoising.py index 653c50b08..ff3a0f181 100644 --- a/fastvideo/pipelines/stages/matrixgame2_denoising.py +++ b/fastvideo/pipelines/stages/matrixgame2_denoising.py @@ -221,7 +221,6 @@ class MatrixGame2CausalDenoisingStage(DenoisingStage): latents[:, :, start_index:start_index + current_num_frames, :, :] = current_latents - # Update KV caches with clean context self._update_context_cache( current_latents=current_latents, batch=batch, @@ -362,7 +361,6 @@ class MatrixGame2CausalDenoisingStage(DenoisingStage): noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4) for i, t_cur in enumerate(timesteps): - # Stop if interrupted if getattr(self, "interrupt", False): break if ctx.boundary_timestep is not None and t_cur < ctx.boundary_timestep: @@ -380,7 +378,6 @@ class MatrixGame2CausalDenoisingStage(DenoisingStage): # t_expand needs to be [batch * frames] to match flattened pred_noise/noise_latents t_expand = t_cur.repeat(latent_model_input.shape[0] * current_num_frames) - # Build attention metadata if VSA is available if vsa_available and self.attn_backend == VideoSparseAttentionBackend: self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls() if self.attn_metadata_builder_cls is not None: @@ -450,7 +447,6 @@ class MatrixGame2CausalDenoisingStage(DenoisingStage): if i < len(timesteps) - 1: next_timestep = timesteps[i + 1] * torch.ones([1], dtype=torch.long, device=pred_video_btchw.device) - # Use custom noise generator if provided (for streaming), else generate if noise_generator is not None: noise = noise_generator(pred_video_btchw.shape, pred_video_btchw.dtype, i) else: @@ -582,7 +578,6 @@ class MatrixGame2CausalDenoisingStage(DenoisingStage): prompt_embeds = batch.prompt_embeds assert torch.isnan(prompt_embeds[0]).sum() == 0 - # Initialize caches kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], dtype=target_dtype, device=latents.device) kv_cache2 = None if boundary_timestep is not None: @@ -603,7 +598,6 @@ class MatrixGame2CausalDenoisingStage(DenoisingStage): dtype=target_dtype, device=latents.device) - # Calculate block sizes if t % self.num_frame_per_block != 0: raise ValueError("num_frames must be divisible by num_frame_per_block for causal denoising") num_blocks = t // self.num_frame_per_block @@ -622,7 +616,6 @@ class MatrixGame2CausalDenoisingStage(DenoisingStage): ) for _ in range(num_denoising_steps - 1) ] - # Create and store context self._streaming_ctx = BlockProcessingContext( batch=batch, block_idx=0, @@ -667,7 +660,6 @@ class MatrixGame2CausalDenoisingStage(DenoisingStage): current_latents = latents[:, :, start_index:start_index + current_num_frames, :, :] - # Update batch with new actions for this block if keyboard_action is not None or mouse_action is not None: vae_ratio = 4 start_frame = 0 if start_index == 0 else 1 + vae_ratio * (start_index - 1) @@ -681,7 +673,6 @@ class MatrixGame2CausalDenoisingStage(DenoisingStage): action_kwargs = self._prepare_action_kwargs(batch, start_index, current_num_frames) - # Create noise generator that uses pre-allocated noise pool def streaming_noise_generator(shape: tuple, dtype: torch.dtype, step_idx: int) -> torch.Tensor: if ctx.noise_pool is not None and step_idx < len(ctx.noise_pool): return ctx.noise_pool[step_idx][:, :shape[1], :, :, :].to(latents.device) @@ -706,7 +697,6 @@ class MatrixGame2CausalDenoisingStage(DenoisingStage): latents[:, :, start_index:start_index + current_num_frames, :, :] = current_latents - # Update KV caches with clean context self._update_context_cache( current_latents=current_latents, batch=batch, diff --git a/fastvideo/pipelines/stages/matrixgame3_denoising.py b/fastvideo/pipelines/stages/matrixgame3_denoising.py index c8fcd652c..6c5b1fc19 100644 --- a/fastvideo/pipelines/stages/matrixgame3_denoising.py +++ b/fastvideo/pipelines/stages/matrixgame3_denoising.py @@ -246,7 +246,6 @@ class MatrixGame3DenoisingStage(DenoisingStage): (x_memory.shape[0], x_memory.shape[2] * x_memory.shape[3] * x_memory.shape[4] // 4)) with self.progress_bar(total=len(timesteps)) as progress_bar: for timestep in timesteps: - # Stop if interrupted if getattr(self, "interrupt", False): break latent_model_input = current_latents diff --git a/fastvideo/pipelines/stages/sd35_conditioning.py b/fastvideo/pipelines/stages/sd35_conditioning.py index d2033e89b..f76982167 100644 --- a/fastvideo/pipelines/stages/sd35_conditioning.py +++ b/fastvideo/pipelines/stages/sd35_conditioning.py @@ -241,7 +241,6 @@ class SD35DenoisingStage(PipelineStage): {"generator": batch.generator[0] if isinstance(batch.generator, list) else batch.generator}) for t in timesteps: - # Stop if interrupted if getattr(self, "interrupt", False): break latents_4d = latents.squeeze(2) diff --git a/fastvideo/pipelines/stages/sr_denoising.py b/fastvideo/pipelines/stages/sr_denoising.py index 2f29ad68e..0f32be765 100644 --- a/fastvideo/pipelines/stages/sr_denoising.py +++ b/fastvideo/pipelines/stages/sr_denoising.py @@ -111,9 +111,6 @@ class SRDenoisingStage(PipelineStage): pipeline.add_module("upsampler", self.upsampler) fastvideo_args.model_loaded["upsampler"] = True - # Setup precision and autocast settings - # TODO(will): make the precision configurable for inference - # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision] target_dtype = torch.bfloat16 autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast @@ -146,7 +143,6 @@ class SRDenoisingStage(PipelineStage): }, ) - # Get latents and embeddings prompt_embeds = batch.prompt_embeds assert not torch.isnan(prompt_embeds[0]).any(), "prompt_embeds contains nan" @@ -174,10 +170,8 @@ class SRDenoisingStage(PipelineStage): latent_model_input = latents.to(target_dtype) assert latent_model_input.shape[0] == 1, "only support batch size 1" - # Run denoising loop with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): - # Skip if interrupted if hasattr(self, 'interrupt') and self.interrupt: continue @@ -198,7 +192,6 @@ class SRDenoisingStage(PipelineStage): latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) - # Prepare inputs for transformer guidance_expand = (torch.tensor( [fastvideo_args.pipeline_config.embedded_cfg_scale] * latent_model_input.shape[0], dtype=torch.float32, @@ -212,7 +205,6 @@ class SRDenoisingStage(PipelineStage): if self.attn_metadata_builder_cls is not None: self.attn_metadata_builder = self.attn_metadata_builder_cls() - # TODO(will): clean this up attn_metadata = self.attn_metadata_builder.build( # type: ignore current_timestep=i, # type: ignore raw_latent_shape=batch.raw_latent_shape[2:5], # type: ignore @@ -228,7 +220,6 @@ class SRDenoisingStage(PipelineStage): self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls() if self.attn_metadata_builder_cls is not None: self.attn_metadata_builder = self.attn_metadata_builder_cls() - # Prepare V-MoBA parameters from config moba_params = fastvideo_args.moba_config.copy() moba_params.update({ "current_timestep": i, @@ -242,18 +233,12 @@ class SRDenoisingStage(PipelineStage): attn_metadata = None else: attn_metadata = None - # TODO(will): finalize the interface. vLLM uses this to - # support torch dynamo compilation. They pass in - # attn_metadata, vllm_config, and num_tokens. We can pass in - # fastvideo_args or training_args, and attn_metadata. batch.is_cfg_negative = False with set_forward_context( current_timestep=i, attn_metadata=attn_metadata, forward_batch=batch, - # fastvideo_args=fastvideo_args ): - # Run transformer noise_pred = self.transformer(latent_model_input, prompt_embeds, t_expand, @@ -262,15 +247,12 @@ class SRDenoisingStage(PipelineStage): **pos_cond_kwargs, **image_kwargs) - # Compute the previous noisy sample latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] - # Update progress bar if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 and progress_bar is not None): progress_bar.update() - # Update batch with final latents batch.latents = latents # deallocate transformer if on mps diff --git a/fastvideo/pipelines/stages/text_encoding.py b/fastvideo/pipelines/stages/text_encoding.py index addccf4ba..343d3f684 100644 --- a/fastvideo/pipelines/stages/text_encoding.py +++ b/fastvideo/pipelines/stages/text_encoding.py @@ -61,7 +61,6 @@ class TextEncodingStage(PipelineStage): assert len(self.tokenizers) == len(self.text_encoders) assert len(self.text_encoders) == len(fastvideo_args.pipeline_config.text_encoder_configs) - # Skip encoding if precomputed prompt_embeds were provided if batch.prompt_embeds is not None and len(batch.prompt_embeds) > 0: return batch @@ -188,16 +187,13 @@ class TextEncodingStage(PipelineStage): indices = [encoder_index] else: indices = list(encoder_index) - # validate range num_encoders = len(self.text_encoders) for idx in indices: if idx < 0 or idx >= num_encoders: raise IndexError(f"encoder index {idx} out of range [0, {num_encoders-1}]") - # Validate indices are within range num_encoders = len(self.text_encoders) - # Normalize input to list[str] assert isinstance(text, str | list) if isinstance(text, str): texts: list[str] = [text] diff --git a/fastvideo/pipelines/stages/timestep_preparation.py b/fastvideo/pipelines/stages/timestep_preparation.py index 91441b34d..adb2c9567 100644 --- a/fastvideo/pipelines/stages/timestep_preparation.py +++ b/fastvideo/pipelines/stages/timestep_preparation.py @@ -53,12 +53,10 @@ class TimestepPreparationStage(PipelineStage): sigmas = batch.sigmas n_tokens = batch.n_tokens - # Prepare extra kwargs for set_timesteps extra_set_timesteps_kwargs = {} if n_tokens is not None and "n_tokens" in inspect.signature(scheduler.set_timesteps).parameters: extra_set_timesteps_kwargs["n_tokens"] = n_tokens - # Handle custom timesteps or sigmas if timesteps is not None and sigmas is not None: raise ValueError( "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") @@ -85,7 +83,6 @@ class TimestepPreparationStage(PipelineStage): scheduler.set_timesteps(num_inference_steps, device=device, **extra_set_timesteps_kwargs) timesteps = scheduler.timesteps - # Update batch with prepared timesteps batch.timesteps = timesteps return batch diff --git a/fastvideo/pipelines/stages/validators.py b/fastvideo/pipelines/stages/validators.py index 252ab7455..f51cccc67 100644 --- a/fastvideo/pipelines/stages/validators.py +++ b/fastvideo/pipelines/stages/validators.py @@ -363,7 +363,6 @@ class VerificationResult: failures = [] all_passed = True - # Apply all validators and collect detailed failure info for validator in validators: try: passed = validator(value) @@ -394,7 +393,6 @@ class VerificationResult: expected = None error_msg = None - # Handle common validator patterns if hasattr(validator, '__closure__') and validator.__closure__: # This is likely a closure (like our helper functions) if 'dims' in validator_name or 'with_dims' in str(validator): @@ -405,7 +403,6 @@ class VerificationResult: elif 'divisible' in str(validator): expected = f"integer divisible by {validator.__closure__[0].cell_contents}" - # Handle specific validator types and check for NaN values if validator_name == 'is_tensor': expected = "torch.Tensor without NaN values" if isinstance(value, torch.Tensor) and torch.isnan(value).any().item(): diff --git a/fastvideo/platforms/__init__.py b/fastvideo/platforms/__init__.py index dc957064d..b687cd3e3 100644 --- a/fastvideo/platforms/__init__.py +++ b/fastvideo/platforms/__init__.py @@ -7,7 +7,6 @@ import traceback from typing import TYPE_CHECKING from fastvideo.logger import init_logger -# imported by other files, do not remove from fastvideo.platforms.interface import AttentionBackendEnum # noqa: F401 from fastvideo.platforms.interface import Platform, PlatformEnum from fastvideo.utils import resolve_obj_by_qualname @@ -151,8 +150,6 @@ builtin_platform_plugins = { def resolve_current_platform_cls_qualname() -> str: - # TODO(will): if we need to support other platforms, we should consider if - # vLLM's plugin architecture is suitable for our needs. # Try MPS first on macOS platform_cls_qualname = mps_platform_plugin() diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index 78c6cdc04..d00b2f0ae 100644 --- a/fastvideo/platforms/cuda.py +++ b/fastvideo/platforms/cuda.py @@ -168,8 +168,6 @@ class CudaPlatformBase(Platform): @classmethod def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, head_size: int, dtype: torch.dtype) -> str: - # TODO(will): maybe come up with a more general interface for local attention - # if distributed is False, we always try to use Flash attn logger.info("Trying FASTVIDEO_ATTENTION_BACKEND=%s", envs.FASTVIDEO_ATTENTION_BACKEND.get()) logger.info("Selected backend: %s", selected_backend) diff --git a/fastvideo/platforms/mps.py b/fastvideo/platforms/mps.py index 6f8df1ff1..7fc0c9f2b 100644 --- a/fastvideo/platforms/mps.py +++ b/fastvideo/platforms/mps.py @@ -62,7 +62,6 @@ class MpsPlatform(Platform): @classmethod def get_device_communicator_cls(cls) -> str: - # Use base communicator for MPS return "fastvideo.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" @classmethod diff --git a/fastvideo/profiler.py b/fastvideo/profiler.py index bcfc9ab84..1f1859ef9 100644 --- a/fastvideo/profiler.py +++ b/fastvideo/profiler.py @@ -134,10 +134,6 @@ register_profiler_region( name="profiler_region_model_loading", description="Module/model loading during pipeline initialization.", ) -# register_profiler_region( -# name="profiler_region_inference_pre_denoising", -# description="Pre-denoising inference steps (conditioning, preprocessing).", -# ) register_profiler_region( name="profiler_region_inference_denoising", description="The main inference denoising loop.", diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 0a0d18361..240508020 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -572,7 +572,6 @@ def _register_configs() -> None: model_family="flux2", default_preset="flux2_klein_4b", ) - # Flux2 (full, Mistral3 text encoder, embedded guidance) register_configs( sampling_param_cls=None, pipeline_config_cls=Flux2PipelineConfig, @@ -645,7 +644,6 @@ def _register_configs() -> None: default_preset="hunyuan15_sr_1080p", ) - # Hunyuan (excludes hyworld, versioned models, and removed GameCraft checkpoints) register_configs( sampling_param_cls=None, pipeline_config_cls=HunyuanConfig, diff --git a/fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py b/fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py index 93ecf8dd3..847b1a171 100644 --- a/fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py +++ b/fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py @@ -128,7 +128,6 @@ def test_grad_and_inference_steps_share_one_builder_safely(): xg = torch.randn(1, meta_grad.total_seq_length, 2, 8, requires_grad=True) impl.tile(xg, meta_grad) - # Then an inference step: the buffer is allocated and reused as usual. meta_inf, _ = _build(builder, step=1) x = torch.randn(1, meta_inf.total_seq_length, 2, 8) with torch.no_grad(): diff --git a/fastvideo/tests/dataset/test_parquet_io.py b/fastvideo/tests/dataset/test_parquet_io.py index 2518130e1..9992bc125 100644 --- a/fastvideo/tests/dataset/test_parquet_io.py +++ b/fastvideo/tests/dataset/test_parquet_io.py @@ -55,7 +55,6 @@ def test_writer_flush_and_remainder(tmp_path: Path): total_rows = sum(pq.read_table(str(f)).num_rows for f in files) assert total_rows == 20 - # Append remainder to complete another chunk extra = records_to_table([{"id": str(i)} for i in range(5)], schema) writer.append_table(extra) written2 = writer.flush(num_workers=1) @@ -99,7 +98,6 @@ def test_writer_parallel_workers(tmp_path: Path): written = writer.flush(num_workers=2) assert written == 40 - # Ensure files exist under worker subdirs worker_dirs = [p for p in out_dir.iterdir() if p.is_dir() and p.name.startswith("worker_")] assert len(worker_dirs) >= 1 files = sorted(out_dir.rglob("*.parquet")) diff --git a/fastvideo/tests/encoders/test_clip_encoder.py b/fastvideo/tests/encoders/test_clip_encoder.py index 95065d235..fb454a03a 100644 --- a/fastvideo/tests/encoders/test_clip_encoder.py +++ b/fastvideo/tests/encoders/test_clip_encoder.py @@ -43,13 +43,11 @@ def test_clip_encoder(): logger.info("Loading models from %s", args.model_path) - # config = json.load(open(os.path.join(model_path, "config.json"))) hf_config = AutoConfig.from_pretrained(TEXT_ENCODER_PATH) print(hf_config) print(hf_config.use_return_dict) - # Load HuggingFace implementation model1 = CLIPTextModel.from_pretrained(TEXT_ENCODER_PATH).to(torch.float16).to(device).eval() from fastvideo.models.loader.component_loader import TextEncoderLoader @@ -57,9 +55,6 @@ def test_clip_encoder(): loader = TextEncoderLoader() model2 = loader.load(TEXT_ENCODER_PATH, args) - # Load the HuggingFace implementation directly - # model2 = CLIPTextModel(hf_config) - # model2 = model2.to(torch.float16) model2.eval() # Sanity check weights between the two models @@ -67,7 +62,6 @@ def test_clip_encoder(): params1 = dict(model1.named_parameters()) params2 = dict(model2.named_parameters()) - # Check number of parameters logger.info("Model1 has %d parameters", len(params1)) logger.info("Model2 has %d parameters", len(params2)) @@ -85,7 +79,6 @@ def test_clip_encoder(): assert_close(param1, param2, atol=1e-4, rtol=1e-4) gc.collect() torch.cuda.empty_cache() - # Load tokenizer tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH) # Test with some sample prompts @@ -100,24 +93,19 @@ def test_clip_encoder(): # Tokenize the prompt tokens = tokenizer(prompt, padding="max_length", max_length=77, truncation=True, return_tensors="pt").to(device) - # Get embeddings from our implementation outputs1 = model1(input_ids=tokens.input_ids, output_hidden_states=True) logger.info("Testing model2") print("--------------------------------") - # Get embeddings from HuggingFace implementation with set_forward_context(current_timestep=0, attn_metadata=None): outputs2 = model2( input_ids=tokens.input_ids, - # attention_mask=tokens.attention_mask, output_hidden_states=True, ) # Compare last hidden states last_hidden_state1 = outputs1.last_hidden_state[tokens.attention_mask == 1] last_hidden_state2 = outputs2.last_hidden_state[tokens.attention_mask == 1] - # print("last_hidden_state1", last_hidden_state1) - # print("last_hidden_state2", last_hidden_state2) assert last_hidden_state1.shape == last_hidden_state2.shape, ( f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}") diff --git a/fastvideo/tests/encoders/test_hyt5_encoder.py b/fastvideo/tests/encoders/test_hyt5_encoder.py index 568fce3c5..e5510863d 100644 --- a/fastvideo/tests/encoders/test_hyt5_encoder.py +++ b/fastvideo/tests/encoders/test_hyt5_encoder.py @@ -34,7 +34,6 @@ def t5_model_paths_and_config(): @pytest.mark.usefixtures("distributed_setup") def test_t5_encoder(t5_model_paths_and_config): - # Initialize the two model implementations text_encoder_path, tokenizer_path, pipeline_config = t5_model_paths_and_config hf_config = AutoConfig.from_pretrained(text_encoder_path) print(hf_config) @@ -56,11 +55,9 @@ def test_t5_encoder(t5_model_paths_and_config): params1 = dict(model1.named_parameters()) params2 = dict(model2.named_parameters()) - # Check number of parameters logger.info("Model1 has %s parameters", len(params1)) logger.info("Model2 has %s parameters", len(params2)) - # check if embed_tokens are the same weights = ["encoder.block.{}.layer.0.layer_norm.weight", \ "encoder.block.{}.layer.0.SelfAttention.o.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_0.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_1.weight",\ "encoder.block.{}.layer.1.DenseReluDense.wo.weight", \ @@ -92,15 +89,10 @@ def test_t5_encoder(t5_model_paths_and_config): add_special_tokens=True, return_tensors="pt").to(device) - # Get outputs from HuggingFace implementation - # filter out padding input_ids - # tokens.input_ids = tokens.input_ids[tokens.attention_mask==1] - # tokens.attention_mask = tokens.attention_mask[tokens.attention_mask==1] outputs1 = model1(input_ids=tokens.input_ids, attention_mask=tokens.attention_mask.float())[0] print("--------------------------------") logger.info("Testing model2") - # Get outputs from our implementation with set_forward_context(current_timestep=0, attn_metadata=None): outputs2 = model2( input_ids=tokens.input_ids, diff --git a/fastvideo/tests/encoders/test_llama_encoder.py b/fastvideo/tests/encoders/test_llama_encoder.py index 6b940160f..d6a709750 100644 --- a/fastvideo/tests/encoders/test_llama_encoder.py +++ b/fastvideo/tests/encoders/test_llama_encoder.py @@ -43,19 +43,15 @@ def test_llama_encoder(): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - # Initialize the two model implementations logger.info("Loading models from %s", args.model_path) hf_config = AutoConfig.from_pretrained(TEXT_ENCODER_PATH) print(hf_config) - # Load HuggingFace implementation model1 = LlamaModel.from_pretrained(TEXT_ENCODER_PATH).to(torch.float16).to(device).eval() loader = TextEncoderLoader() device = torch.device("cuda:0") model2 = loader.load(TEXT_ENCODER_PATH, args) - # Convert to float16 and move to device - # model2 = model2.to(torch.float16) model2.eval() # Sanity check weights between the two models @@ -63,11 +59,9 @@ def test_llama_encoder(): params1 = dict(model1.named_parameters()) params2 = dict(model2.named_parameters()) - # Check number of parameters logger.info("Model1 has %d parameters", len(params1)) logger.info("Model2 has %d parameters", len(params2)) - # check if embed_tokens are the same device = model1.embed_tokens.weight.device assert torch.allclose( model1.embed_tokens.weight, @@ -82,7 +76,6 @@ def test_llama_encoder(): p1 = params1[name1] p2 = params2[name2] if "gate_up" in name2: - # print("skipping gate_up") continue p1 = p1.to_local().to(device) if isinstance(p1, DTensor) else p1.to(device) p2 = p2.to_local().to(device) if isinstance(p2, DTensor) else p2.to(device) @@ -122,15 +115,10 @@ def test_llama_encoder(): tokens = tokenizer(prompt, padding="max_length", max_length=512, truncation=True, return_tensors="pt").to(device) - # Get outputs from our implementation - # filter out padding input_ids - # tokens.input_ids = tokens.input_ids[tokens.attention_mask==1] - # tokens.attention_mask = tokens.attention_mask[tokens.attention_mask==1] outputs1 = model1(input_ids=tokens.input_ids, output_hidden_states=True) print("--------------------------------") logger.info("Testing model2") - # Get outputs from HuggingFace implementation with set_forward_context(current_timestep=0, attn_metadata=None): outputs2 = model2(input_ids=tokens.input_ids, attention_mask=tokens.attention_mask, diff --git a/fastvideo/tests/encoders/test_qwen2_5_encoder.py b/fastvideo/tests/encoders/test_qwen2_5_encoder.py index b6a09bb5d..bd4869381 100644 --- a/fastvideo/tests/encoders/test_qwen2_5_encoder.py +++ b/fastvideo/tests/encoders/test_qwen2_5_encoder.py @@ -43,11 +43,9 @@ def test_qwen2_5_encoder(qwen_model_path_and_config): logger.info(f"Using precision: {precision_str}") - # Load HF model (Base model) model1 = Qwen2_5_VLTextModel.from_pretrained(text_encoder_path).to(precision).to(device).eval() tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) - # Load FastVideo model args = FastVideoArgs(model_path=text_encoder_path, pipeline_config=pipeline_config, pin_cpu_memory=False) loader = TextEncoderLoader() @@ -81,10 +79,8 @@ def test_qwen2_5_encoder(qwen_model_path_and_config): p2 = params2[name2] p2 = (p2.to_local() if isinstance(p2, DTensor) else p2).to(p1) - # Check shape assert p1.shape == p2.shape, f"Shape mismatch for {w}: {p1.shape} vs {p2.shape}" - # Check values assert_close(p1, p2, atol=1e-7, rtol=1e-7, msg=f"Weight mismatch for {w}") # Test with sample prompts @@ -113,7 +109,6 @@ def test_qwen2_5_encoder(qwen_model_path_and_config): # Compare # Filter padding for comparison if needed, but here we just check raw output matching - # Check shapes assert outputs1.shape == outputs2.shape, f"Output shape mismatch: {outputs1.shape} vs {outputs2.shape}" diff = torch.abs(outputs1 - outputs2) diff --git a/fastvideo/tests/encoders/test_siglip_encoder.py b/fastvideo/tests/encoders/test_siglip_encoder.py index 3900e6292..86cdaa498 100644 --- a/fastvideo/tests/encoders/test_siglip_encoder.py +++ b/fastvideo/tests/encoders/test_siglip_encoder.py @@ -41,19 +41,15 @@ def test_siglip_encoder(): logger.info("Loading SigLIP models from %s/%s", MODEL_ID, IMAGE_ENCODER_SUBFOLDER) - # Load HuggingFace implementation hf_model = HFSiglipVisionModel.from_pretrained(MODEL_ID, subfolder=IMAGE_ENCODER_SUBFOLDER).to( torch.float16).to(device).eval() - # Load the config from Hugging Face and extract vision_config config_path = hf_hub_download(repo_id=MODEL_ID, filename=f"{IMAGE_ENCODER_SUBFOLDER}/config.json") with open(config_path) as f: full_config = json.load(f) - # Get vision config from the full config vision_config_dict = full_config.get("vision_config", full_config) - # Create FastVideo config with vision-specific settings arch_config = SiglipVisionArchConfig( hidden_size=vision_config_dict.get("hidden_size", 1152), image_size=vision_config_dict.get("image_size", 384), @@ -65,11 +61,9 @@ def test_siglip_encoder(): config = SiglipVisionConfig(arch_config=arch_config) - # Create FastVideo model from fastvideo.models.encoders.siglip import SiglipVisionModel fv_model = SiglipVisionModel(config).to(torch.float16).to(device) - # Load weights from safetensors via Hugging Face weights_path = hf_hub_download(repo_id=MODEL_ID, filename=f"{IMAGE_ENCODER_SUBFOLDER}/model.safetensors") state_dict = load_file(weights_path) # Filter to only vision_model weights (keep the vision_model. prefix) @@ -110,16 +104,13 @@ def test_siglip_encoder(): batch_size = 2 image_size = fv_model.config.arch_config.image_size - # Create random pixel values pixel_values = torch.randn(batch_size, 3, image_size, image_size).to(device).to(torch.float16) logger.info("Testing SigLIP encoder with random pixel values of shape %s", pixel_values.shape) with torch.no_grad(): - # Get embeddings from HuggingFace implementation hf_outputs = hf_model(pixel_values=pixel_values) - # Get embeddings from FastVideo implementation with set_forward_context(current_timestep=0, attn_metadata=None): fv_outputs = fv_model(pixel_values=pixel_values) diff --git a/fastvideo/tests/encoders/test_t5_encoder.py b/fastvideo/tests/encoders/test_t5_encoder.py index 01d40348c..34e69706f 100644 --- a/fastvideo/tests/encoders/test_t5_encoder.py +++ b/fastvideo/tests/encoders/test_t5_encoder.py @@ -46,7 +46,6 @@ def t5_large_model_paths_and_config(): @pytest.mark.usefixtures("distributed_setup") def test_t5_encoder(t5_model_paths_and_config): - # Initialize the two model implementations text_encoder_path, tokenizer_path, pipeline_config = t5_model_paths_and_config hf_config = AutoConfig.from_pretrained(text_encoder_path) print(hf_config) @@ -68,11 +67,9 @@ def test_t5_encoder(t5_model_paths_and_config): params1 = dict(model1.named_parameters()) params2 = dict(model2.named_parameters()) - # Check number of parameters logger.info("Model1 has %s parameters", len(params1)) logger.info("Model2 has %s parameters", len(params2)) - # check if embed_tokens are the same weights = ["encoder.block.{}.layer.0.layer_norm.weight", "encoder.block.{}.layer.0.SelfAttention.relative_attention_bias.weight", \ "encoder.block.{}.layer.0.SelfAttention.o.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_0.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_1.weight",\ "encoder.block.{}.layer.1.DenseReluDense.wo.weight", \ @@ -100,17 +97,12 @@ def test_t5_encoder(t5_model_paths_and_config): tokens = tokenizer(prompt, padding="max_length", max_length=512, truncation=True, return_tensors="pt").to(device) - # Get outputs from HuggingFace implementation - # filter out padding input_ids - # tokens.input_ids = tokens.input_ids[tokens.attention_mask==1] - # tokens.attention_mask = tokens.attention_mask[tokens.attention_mask==1] outputs1 = model1(input_ids=tokens.input_ids, attention_mask=tokens.attention_mask, output_hidden_states=True).last_hidden_state print("--------------------------------") logger.info("Testing model2") - # Get outputs from our implementation with set_forward_context(current_timestep=0, attn_metadata=None): outputs2 = model2( input_ids=tokens.input_ids, @@ -126,7 +118,6 @@ def test_t5_encoder(t5_model_paths_and_config): @pytest.mark.usefixtures("distributed_setup") def test_t5_large_encoder(t5_large_model_paths_and_config): - # Initialize the two model implementations text_encoder_path, tokenizer_path, pipeline_config = t5_large_model_paths_and_config hf_config = AutoConfig.from_pretrained(text_encoder_path) print(hf_config) @@ -148,7 +139,6 @@ def test_t5_large_encoder(t5_large_model_paths_and_config): params1 = dict(model1.named_parameters()) params2 = dict(model2.named_parameters()) - # Check number of parameters logger.info("Model1 has %s parameters", len(params1)) logger.info("Model2 has %s parameters", len(params2)) @@ -161,7 +151,6 @@ def test_t5_large_encoder(t5_large_model_paths_and_config): for name in sorted(params2.keys()): logger.info(" %s: %s", name, params2[name].shape) - #check if embed_tokens are the same weights = ["encoder.block.{}.layer.0.layer_norm.weight", "encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight", \ "encoder.block.{}.layer.0.SelfAttention.o.weight", "encoder.block.{}.layer.1.DenseReluDense.wi.weight", \ "encoder.block.{}.layer.1.DenseReluDense.wo.weight", "encoder.final_layer_norm.weight"] @@ -188,17 +177,12 @@ def test_t5_large_encoder(t5_large_model_paths_and_config): tokens = tokenizer(prompt, padding="max_length", max_length=512, truncation=True, return_tensors="pt").to(device) - # Get outputs from HuggingFace implementation - # filter out padding input_ids - # tokens.input_ids = tokens.input_ids[tokens.attention_mask==1] - # tokens.attention_mask = tokens.attention_mask[tokens.attention_mask==1] outputs1 = model1(input_ids=tokens.input_ids, attention_mask=tokens.attention_mask, output_hidden_states=True).last_hidden_state print("--------------------------------") logger.info("Testing model2 with T5LargeConfig") - # Get outputs from our implementation with set_forward_context(current_timestep=0, attn_metadata=None): outputs2 = model2( input_ids=tokens.input_ids, diff --git a/fastvideo/tests/entrypoints/streaming/test_prompt_providers.py b/fastvideo/tests/entrypoints/streaming/test_prompt_providers.py index dc5349bea..623db0204 100644 --- a/fastvideo/tests/entrypoints/streaming/test_prompt_providers.py +++ b/fastvideo/tests/entrypoints/streaming/test_prompt_providers.py @@ -89,9 +89,6 @@ def _install_fake_httpx( return captured -# ---------------------------------------------------------------------- -# Cerebras -# ---------------------------------------------------------------------- class TestCerebrasProvider: @@ -141,7 +138,6 @@ class TestCerebrasProvider: assert result.content == "enhanced" assert result.provider == "cerebras" assert result.model == "gpt-oss-120b" - # Verify the HTTP request was shaped correctly. body = captured[0]["json"] assert body["model"] == "gpt-oss-120b" assert body["messages"] == [{"role": "user", "content": "a fox"}] @@ -251,9 +247,6 @@ class TestCerebrasProvider: asyncio.run(provider.complete(LLMRequest(messages=[], model="m"))) -# ---------------------------------------------------------------------- -# Groq -# ---------------------------------------------------------------------- class TestGroqProvider: diff --git a/fastvideo/tests/entrypoints/streaming/test_server.py b/fastvideo/tests/entrypoints/streaming/test_server.py index 6e9280abb..c98e32cd5 100644 --- a/fastvideo/tests/entrypoints/streaming/test_server.py +++ b/fastvideo/tests/entrypoints/streaming/test_server.py @@ -259,9 +259,6 @@ class TestContinuationStatePersistence: assert snap["state"]["payload"]["source_prompt"] == "a cat" -# ---------------------------------------------------------------------- -# Helpers -# ---------------------------------------------------------------------- @contextlib.contextmanager @@ -296,7 +293,6 @@ def _drain_until(ws, target_type: str) -> dict[str, Any]: parsed = json.loads(msg["text"]) if parsed.get("type") == target_type: return parsed - # skip binary / other def _as_json(msg: dict[str, Any]) -> dict[str, Any] | None: diff --git a/fastvideo/tests/entrypoints/streaming/test_session.py b/fastvideo/tests/entrypoints/streaming/test_session.py index dd0748df1..e81f1ab31 100644 --- a/fastvideo/tests/entrypoints/streaming/test_session.py +++ b/fastvideo/tests/entrypoints/streaming/test_session.py @@ -95,7 +95,6 @@ class TestSessionManager: s = mgr.create() mgr.close(s.id) assert len(mgr) == 0 - # Now can create again. mgr.create() def test_reap_timed_out_flags_idle_sessions(self): diff --git a/fastvideo/tests/hooks/test_hooks_api.py b/fastvideo/tests/hooks/test_hooks_api.py index 4f6b277cf..94352d7e3 100644 --- a/fastvideo/tests/hooks/test_hooks_api.py +++ b/fastvideo/tests/hooks/test_hooks_api.py @@ -26,13 +26,10 @@ class EventHook(ForwardHook): def test_hook_execution_order(): """Test that hooks are executed in the correct order: LIFO for pre-hooks, FIFO for post-hooks.""" - # Create a simple model model = nn.Linear(10, 20) - # Create event list to track hook execution order events = [] - # Create and push hooks in order: A then B manager = ModuleHookManager.get_from_or_default(model) diff --git a/fastvideo/tests/hooks/test_layerwise_offload.py b/fastvideo/tests/hooks/test_layerwise_offload.py index f1505029e..86a9b3af4 100644 --- a/fastvideo/tests/hooks/test_layerwise_offload.py +++ b/fastvideo/tests/hooks/test_layerwise_offload.py @@ -53,10 +53,8 @@ def test_layerwise_offload_basic(): seq_len = 16 num_blocks = 4 - # Create model model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32).to(device) - # Get reference output without offloading input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device) with torch.no_grad(): reference_output = model(input_tensor.clone()) @@ -64,17 +62,14 @@ def test_layerwise_offload_basic(): # Enable layerwise offloading enable_layerwise_offload(model) - # Verify parameters are offloaded to CPU for block in model.blocks: for param in block.parameters(): # Parameters should be placeholder tensors (empty) assert param.numel() == 0, ("Parameters should be offloaded (empty tensors)") - # Run forward pass with offloading with torch.no_grad(): offloaded_output = model(input_tensor.clone()) - # Check output correctness assert torch.allclose(reference_output, offloaded_output, rtol=1e-4, atol=1e-5), "Output with offloading should match reference output" @@ -91,10 +86,8 @@ def test_layerwise_offload_bf16(): seq_len = 32 num_blocks = 3 - # Create model model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.bfloat16).to(device) - # Get reference output without offloading input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device, dtype=torch.bfloat16) with torch.no_grad(): reference_output = model(input_tensor.clone()) @@ -102,7 +95,6 @@ def test_layerwise_offload_bf16(): # Enable layerwise offloading enable_layerwise_offload(model) - # Run forward pass with offloading with torch.no_grad(): offloaded_output = model(input_tensor.clone()) @@ -124,10 +116,8 @@ def test_layerwise_offload_multiple_forward_passes(): num_blocks = 3 num_iterations = 5 - # Create model model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32).to(device) - # Get reference outputs without offloading torch.manual_seed(42) reference_outputs = [] for i in range(num_iterations): @@ -138,14 +128,12 @@ def test_layerwise_offload_multiple_forward_passes(): # Enable layerwise offloading enable_layerwise_offload(model) - # Run multiple forward passes with offloading torch.manual_seed(42) for i in range(num_iterations): input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device) with torch.no_grad(): offloaded_output = model(input_tensor.clone()) - # Check output correctness for each iteration assert torch.allclose(reference_outputs[i], offloaded_output, rtol=1e-4, atol=1e-5), f"Output mismatch in iteration {i}" @@ -159,10 +147,8 @@ def test_layerwise_offload_parameter_integrity(): hidden_size = 128 num_blocks = 2 - # Create model model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32).to(device) - # Store original parameter values per block original_params_per_block = [] for block in model.blocks: block_params = {} @@ -173,18 +159,15 @@ def test_layerwise_offload_parameter_integrity(): # Enable layerwise offloading enable_layerwise_offload(model) - # Create hook managers and verify they exist for block_idx, block in enumerate(model.blocks): manager = ModuleHookManager.get_from(block) assert manager is not None, "Hook manager should be attached to blocks" hook: LayerwiseOffloadHook | None = manager.get_forward_hook("LayerwiseOffloadHook") assert hook is not None, "LayerwiseOffloadHook should be registered" - # Verify parameters are stored in CPU state = hook.state assert len(state.cpu_named_parameters) > 0, ("CPU parameters should be stored") - # Verify CPU parameters match original values for this block original_params = original_params_per_block[block_idx] for name, cpu_param in state.cpu_named_parameters.items(): assert name in original_params, (f"CPU parameter {name} not found in original params for block {block_idx}") @@ -227,7 +210,6 @@ def test_layerwise_offload_memory_reduction(): batch_size = 1 seq_len = 64 - # Create model model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32).to(device) # Measure initial GPU memory @@ -240,7 +222,6 @@ def test_layerwise_offload_memory_reduction(): memory_without_offload = torch.cuda.max_memory_allocated() - # Reset model model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32).to(device) # Enable offloading @@ -277,7 +258,6 @@ def test_layerwise_offload_gradient_disabled(): seq_len = 16 num_blocks = 3 - # Create model model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32).to(device) model.eval() @@ -310,7 +290,6 @@ def test_layerwise_offload_different_batch_sizes(): model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32).to(device) - # Get reference model without offloading reference_model = SimpleModelWithModuleList(num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32).to(device) reference_model.load_state_dict(model.state_dict()) @@ -333,7 +312,6 @@ def test_layerwise_offload_different_batch_sizes(): if __name__ == "__main__": - # Run tests manually for debugging if torch.cuda.is_available(): print("Running layerwise offloading tests...") test_layerwise_offload_basic() diff --git a/fastvideo/tests/inference/lora/test_lora_inference_similarity.py b/fastvideo/tests/inference/lora/test_lora_inference_similarity.py index b02788814..5dfaeb615 100644 --- a/fastvideo/tests/inference/lora/test_lora_inference_similarity.py +++ b/fastvideo/tests/inference/lora/test_lora_inference_similarity.py @@ -58,24 +58,6 @@ LORA_CONFIGS = [{ "bad quality video,色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", "ssim_threshold": 0.79 } - # TODO: Add a LoRA with lora_alpha values to test alpha scaling - # - # Context: Some LoRAs use an alpha value that is two times smaller than their rank. - # This fix ensures that LoRA weights are correctly scaled by the alpha/rank ratio when merged. - # - # Issue: Currently, we cannot add a test for LoRA adapters with alpha values because: - # - The existing public LoRAs for Wan-AI/Wan2.1-T2V-1.3B-Diffusers don't store lora_alpha - # - No publicly available LoRA for this model includes lora_alpha tensors in their weights - # - This is why the alpha/rank scaling bug wasn't caught by existing tests - # - # The fix has been validated with: - # - A distilled LoRA that includes alpha values - # - Manual testing shows correct alpha/rank scaling behavior - # - Backward compatibility confirmed with LoRAs without alpha values - # - # Future work: - # - Add a synthetic LoRA test fixture with alpha values when feasible - # - Or wait for public Wan LoRAs with alpha to become available ] MODEL_TO_PARAMS = { diff --git a/fastvideo/tests/lora_extraction/test_lora_extraction.py b/fastvideo/tests/lora_extraction/test_lora_extraction.py index 4ff39980a..897e4a279 100644 --- a/fastvideo/tests/lora_extraction/test_lora_extraction.py +++ b/fastvideo/tests/lora_extraction/test_lora_extraction.py @@ -11,7 +11,6 @@ import torch from fastvideo import VideoGenerator from fastvideo.api import ComponentConfig, EngineConfig, GeneratorConfig, OffloadConfig, ParallelismConfig, PipelineSelection -# Add scripts/lora_extraction to path for imports repo_root = Path(__file__).parents[3] lora_scripts = repo_root / "scripts" / "lora_extraction" sys.path.insert(0, str(lora_scripts)) diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index 8ddbc20b7..3bd0b77eb 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -323,14 +323,7 @@ def run_kernel_tests(): run_test("bash .buildkite/scripts/lanes/kernel_tests.sh") -# @app.function(gpu="H100:1", image=image, timeout=900, secrets=[ci_env_secret]) -# def run_precision_tests_VSA(): -# # VSA correctness is covered by the same file now -# run_test("pytest fastvideo-kernel/tests/test_correctness.py") -# @app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret]) -# def run_precision_tests_vmoba(): -# run_test("pytest fastvideo-kernel/tests/test_vmoba_correctness.py") @app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret]) @@ -359,7 +352,6 @@ def run_unit_test(): run_test("bash .buildkite/scripts/unit_test.sh") -# TODO: David: GPU only used to resolve import time requirement (not needed for this test). Maybe make those imports lazy? @app.function(gpu="L40S:1", image=dreamverse_image, timeout=1800, secrets=[ci_env_secret]) def run_dreamverse_app_tests(): run_test_command(install_command="", diff --git a/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py b/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py index ac154fb9c..8063c940c 100644 --- a/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py +++ b/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py @@ -7,7 +7,6 @@ import shutil import fastvideo.envs as envs from fastvideo.tests.utils import compute_video_ssim_torchvision -# Import the training pipeline sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) NUM_NODES = "1" @@ -32,7 +31,6 @@ LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs")) def download_data(): - # create the data dir if it doesn't exist data_dir = Path(DATA_DIR) print(f"Creating data directory at {data_dir}") @@ -49,7 +47,6 @@ def download_data(): ) print(f"Download completed successfully. Files downloaded to: {result}") - # Verify the download if not LOCAL_RAW_DATA_DIR.exists(): raise RuntimeError(f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist") @@ -65,12 +62,10 @@ def download_data(): def run_preprocessing(): - # remove the local_preprocessed_data_dir if it exists if LOCAL_PREPROCESSED_DATA_DIR.exists(): print(f"Removing local_preprocessed_data_dir: {LOCAL_PREPROCESSED_DATA_DIR}") shutil.rmtree(LOCAL_PREPROCESSED_DATA_DIR) - # Run torchrun command cmd = [ "torchrun", "--nnodes", @@ -227,11 +222,9 @@ def test_e2e_overfit_single_sample(): final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4") print(f"final_validation_video_file: {final_validation_video_file}") - # Ensure both files exist assert os.path.exists(reference_video_file), f"Reference video not found at {reference_video_file}" assert os.path.exists(final_validation_video_file), f"Validation video not found at {final_validation_video_file}" - # Compute SSIM mean_ssim, min_ssim, max_ssim = compute_video_ssim_torchvision( reference_video_file, final_validation_video_file, diff --git a/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py b/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py index 7f93cf90e..b1c0275fb 100644 --- a/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py +++ b/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py @@ -7,7 +7,6 @@ import sys import fastvideo.envs as envs from fastvideo.tests.utils import compute_video_ssim_torchvision -# Import the training pipeline sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) NUM_NODES = "1" @@ -32,7 +31,6 @@ LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs")) def download_data(): - # create the data dir if it doesn't exist data_dir = Path(DATA_DIR) print(f"Creating data directory at {data_dir}") @@ -49,7 +47,6 @@ def download_data(): ) print(f"Download completed successfully. Files downloaded to: {result}") - # Verify the download if not LOCAL_RAW_DATA_DIR.exists(): raise RuntimeError(f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist") @@ -65,7 +62,6 @@ def download_data(): def run_preprocessing(): - # Run torchrun command cmd = [ "torchrun", "--nnodes", @@ -209,11 +205,9 @@ def test_e2e_overfit_single_sample(): final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4") print(f"final_validation_video_file: {final_validation_video_file}") - # Ensure both files exist assert os.path.exists(reference_video_file), f"Reference video not found at {reference_video_file}" assert os.path.exists(final_validation_video_file), f"Validation video not found at {final_validation_video_file}" - # Compute SSIM mean_ssim, min_ssim, max_ssim = compute_video_ssim_torchvision( reference_video_file, final_validation_video_file, diff --git a/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py b/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py index 90a5e03ce..750468c11 100644 --- a/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py +++ b/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py @@ -7,7 +7,6 @@ import sys import fastvideo.envs as envs from fastvideo.tests.utils import compute_video_ssim_torchvision -# Import the training pipeline sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) NUM_NODES = "1" @@ -32,7 +31,6 @@ LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs")) def download_data(): - # create the data dir if it doesn't exist data_dir = Path(DATA_DIR) print(f"Creating data directory at {data_dir}") @@ -49,7 +47,6 @@ def download_data(): ) print(f"Download completed successfully. Files downloaded to: {result}") - # Verify the download if not LOCAL_RAW_DATA_DIR.exists(): raise RuntimeError(f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist") @@ -65,12 +62,10 @@ def download_data(): def run_preprocessing(): - # remove the local_preprocessed_data_dir if it exists if LOCAL_PREPROCESSED_DATA_DIR.exists(): print(f"Removing local_preprocessed_data_dir: {LOCAL_PREPROCESSED_DATA_DIR}") shutil.rmtree(LOCAL_PREPROCESSED_DATA_DIR) - # Run torchrun command cmd = [ "torchrun", "--nnodes", @@ -214,11 +209,9 @@ def test_e2e_overfit_single_sample(): final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4") print(f"final_validation_video_file: {final_validation_video_file}") - # Ensure both files exist assert os.path.exists(reference_video_file), f"Reference video not found at {reference_video_file}" assert os.path.exists(final_validation_video_file), f"Validation video not found at {final_validation_video_file}" - # Compute SSIM mean_ssim, min_ssim, max_ssim = compute_video_ssim_torchvision( reference_video_file, final_validation_video_file, diff --git a/fastvideo/tests/ssim/test_ltx2_similarity.py b/fastvideo/tests/ssim/test_ltx2_similarity.py index 73031a242..e58808256 100644 --- a/fastvideo/tests/ssim/test_ltx2_similarity.py +++ b/fastvideo/tests/ssim/test_ltx2_similarity.py @@ -51,8 +51,6 @@ LTX2_DISTILLED_PARAMS = { "fps": 24, "ltx2_vae_tiling": True, } -# TODO: Regenerate LTX2-Distilled references for the new distilled defaults, -# then remove these historical full-guidance pins. LTX2_DISTILLED_REFERENCE_GUIDANCE_OVERRIDES = { "ltx2_modality_scale_video": 3.0, "ltx2_modality_scale_audio": 3.0, diff --git a/fastvideo/tests/ssim/test_matrixgame2_similarity.py b/fastvideo/tests/ssim/test_matrixgame2_similarity.py index 55f03e96a..f3e7d56dc 100644 --- a/fastvideo/tests/ssim/test_matrixgame2_similarity.py +++ b/fastvideo/tests/ssim/test_matrixgame2_similarity.py @@ -109,7 +109,6 @@ def test_matrixgame2_similarity(prompt, ATTENTION_BACKEND, model_id, env_overrid BASE_PARAMS = params_map[model_id] num_inference_steps = BASE_PARAMS["num_inference_steps"] - # Create action conditions for Matrix-Game 2.0 actions = create_action_presets(BASE_PARAMS["num_frames"], keyboard_dim=BASE_PARAMS["keyboard_dim"], seed=BASE_PARAMS["seed"]) diff --git a/fastvideo/tests/stages/test_kandinsky6_latent_preparation.py b/fastvideo/tests/stages/test_kandinsky6_latent_preparation.py index 49defaf61..69b2d468f 100644 --- a/fastvideo/tests/stages/test_kandinsky6_latent_preparation.py +++ b/fastvideo/tests/stages/test_kandinsky6_latent_preparation.py @@ -65,9 +65,7 @@ def _batch(num_frames: int, *, fps=None) -> types.SimpleNamespace: ) -# -------------------------------------------------------------------------------------------------- # P5: num_frames rounds to the nearest 4k+1 (floor for a non-exact multiple, round up for an exact one) -# -------------------------------------------------------------------------------------------------- @pytest.mark.parametrize("requested,expected", [ @@ -89,9 +87,7 @@ def test_num_frames_rounds_to_nearest_4k_plus_1(cpu_device, requested, expected) assert out.latents.shape[1] == (expected - 1) // 4 + 1 -# -------------------------------------------------------------------------------------------------- # P17: audio latent length follows batch.fps -# -------------------------------------------------------------------------------------------------- def test_audio_latent_length_follows_the_requested_fps_not_the_pipeline_default(cpu_device): diff --git a/fastvideo/tests/stages/test_text_encoding.py b/fastvideo/tests/stages/test_text_encoding.py index 754144b61..09ba8ffc9 100644 --- a/fastvideo/tests/stages/test_text_encoding.py +++ b/fastvideo/tests/stages/test_text_encoding.py @@ -131,7 +131,6 @@ def test_encode_text_selection_and_shapes(): for e in embeds: assert e.shape == (2, hidden) - # with masks embeds2, masks2 = stage.encode_text("a", fastvideo_args, encoder_index=[1], return_attention_mask=True) assert len(embeds2) == 1 and len(masks2) == 1 assert embeds2[0].shape == (1, hidden) diff --git a/fastvideo/tests/train/callbacks/test_ema.py b/fastvideo/tests/train/callbacks/test_ema.py index 3e84079aa..155dcb99d 100644 --- a/fastvideo/tests/train/callbacks/test_ema.py +++ b/fastvideo/tests/train/callbacks/test_ema.py @@ -16,9 +16,6 @@ import torch from fastvideo.train.callbacks.ema import EMACallback -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- class _Student: @@ -54,9 +51,6 @@ def _tiny_transformer(*, fill: float = 0.0) -> torch.nn.Module: return m -# --------------------------------------------------------------------------- -# A. on_train_start -# --------------------------------------------------------------------------- class TestOnTrainStart: @@ -78,9 +72,6 @@ class TestOnTrainStart: cb.on_train_start(_Method(transformer=None), iteration=0) -# --------------------------------------------------------------------------- -# B. on_training_step_end (decay math + start_iter gating) -# --------------------------------------------------------------------------- class TestOnTrainingStepEnd: @@ -156,9 +147,6 @@ class TestOnTrainingStepEnd: assert any(payload.get("ema/decay") == 0.99 and step == 0 for payload, step in tracker.entries) -# --------------------------------------------------------------------------- -# C. ema_context -# --------------------------------------------------------------------------- class TestEmaContext: @@ -190,9 +178,6 @@ class TestEmaContext: assert torch.allclose(transformer.weight, torch.full((2, 4), 9.0)) -# --------------------------------------------------------------------------- -# D. State dict round-trip -# --------------------------------------------------------------------------- class TestStateDict: @@ -212,7 +197,6 @@ class TestStateDict: assert "student_ema" in state assert state["ema_started"] is True - # Build a fresh callback and load. fresh = EMACallback(decay=0.5, start_iter=0) fresh.on_train_start(_Method(_tiny_transformer(fill=0.0)), iteration=0) # Sanity: fresh shadow != saved shadow before load. diff --git a/fastvideo/tests/train/methods/test_pdd_methods.py b/fastvideo/tests/train/methods/test_pdd_methods.py index e67755d4c..2c10c5812 100644 --- a/fastvideo/tests/train/methods/test_pdd_methods.py +++ b/fastvideo/tests/train/methods/test_pdd_methods.py @@ -261,9 +261,6 @@ def _raw_batch(seed: int = 0) -> dict[str, Any]: return {"latent": torch.randn(1, _NUMEL, generator=generator), "prompt": f"prompt-{seed}"} -# ---------------------------------------------------------------------------- -# Config parsing -# ---------------------------------------------------------------------------- def test_parse_pdd_config_defaults_and_validation() -> None: @@ -294,9 +291,6 @@ def test_parse_pdd_config_defaults_and_validation() -> None: parse_pdd_config({"pdd_loss_weight": -1.0}) -# ---------------------------------------------------------------------------- -# Standalone PDD -# ---------------------------------------------------------------------------- def test_pdd_data_latent_step_trains_the_student_from_aligned_block_starts() -> None: @@ -470,9 +464,6 @@ def test_pdd_teacher_guidance_runs_both_forwards_and_mixes_in_decoding_precision assert out.dtype == torch.float32 -# ---------------------------------------------------------------------------- -# DMPDD -# ---------------------------------------------------------------------------- @pytest.mark.parametrize("rollout_mode", ["simulate", "data_latent"]) @@ -494,9 +485,6 @@ def test_dmpdd_rejects_negative_fake_score_weight_decay() -> None: _build(DMPDDMethod, _dmd_method_config(fake_score_weight_decay=-0.01)) -# ---------------------------------------------------------------------------- -# DMPDD -# ---------------------------------------------------------------------------- def test_dmpdd_student_step_uses_the_full_tail_and_carries_one_block() -> None: diff --git a/fastvideo/tests/training/VSA/test_training_loss_VSA.py b/fastvideo/tests/training/VSA/test_training_loss_VSA.py index 78638253a..735bbc199 100644 --- a/fastvideo/tests/training/VSA/test_training_loss_VSA.py +++ b/fastvideo/tests/training/VSA/test_training_loss_VSA.py @@ -6,10 +6,7 @@ import json from huggingface_hub import snapshot_download import torch -# Force VSA to use Triton implementation even on H100 / when CUDA extension is available -# os.environ["FASTVIDEO_KERNEL_VSA_FORCE_TRITON"] = "1" -# Import the training pipeline sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) import fastvideo.envs as envs from fastvideo.training.runner import main @@ -23,10 +20,8 @@ NUM_GPUS_PER_NODE = "2" def run_worker(): """Worker function that will be run on each GPU""" - # Create and populate args parser = build_parser() - # Set the arguments as they are in finetune_v1_test.sh args = parser.parse_args([ "--pipeline_class", "WanTrainingPipeline", "--pipeline_module", "fastvideo.training.wan_training_pipeline", @@ -47,7 +42,6 @@ def run_worker(): "1", "--VSA_sparsity", "0.9" ]) - # Call the main training function main(args) @@ -62,10 +56,8 @@ def test_distributed_training(): repo_type="dataset", local_dir_use_symlinks=False) - # Get the current file path current_file = Path(__file__).resolve() - # Run torchrun command cmd = ["torchrun", "--nnodes", NUM_NODES, "--nproc_per_node", NUM_GPUS_PER_NODE, str(current_file)] # The torchrun workers select VSA when the training pipeline initializes attention. diff --git a/fastvideo/tests/training/Vanilla/mfu_calculation.py b/fastvideo/tests/training/Vanilla/mfu_calculation.py index 5b7f6e00a..7235895ba 100644 --- a/fastvideo/tests/training/Vanilla/mfu_calculation.py +++ b/fastvideo/tests/training/Vanilla/mfu_calculation.py @@ -2,7 +2,6 @@ import os import sys from pathlib import Path -# Set Python path to current folder current_dir = str(Path(__file__).parent.parent.parent.parent.parent) if current_dir not in sys.path: sys.path.insert(0, current_dir) @@ -12,7 +11,6 @@ import torch import json from huggingface_hub import snapshot_download import fastvideo.envs as envs -# Import the training pipeline from fastvideo.training.runner import main from fastvideo.utils import build_parser @@ -35,10 +33,8 @@ envs.setdefault_external("MASTER_PORT", MASTER_PORT) def run_worker(): """Worker function that will be run on each GPU""" - # Create and populate args parser = build_parser() - # Set the arguments as they are in finetune_t2v.sh args = parser.parse_args([ "--pipeline_class", "WanTrainingPipeline", "--pipeline_module", "fastvideo.training.wan_training_pipeline", @@ -121,7 +117,6 @@ def run_worker(): "1" ]) - # Call the main training function main(args) @@ -136,10 +131,8 @@ def test_distributed_training(): repo_type="dataset", local_dir_use_symlinks=False) - # Get the current file path current_file = Path(__file__).resolve() - # Run torchrun command cmd = [ "torchrun", "--nnodes", NUM_NODES, "--nproc_per_node", NUM_GPUS_PER_NODE, "--master_port", MASTER_PORT, str(current_file) @@ -165,10 +158,8 @@ def test_distributed_training(): with summary_file.open() as f: wandb_summary = json.load(f) - # Calculate and print MFU metrics device_name = torch.cuda.get_device_name() try: - # Get actual values from training run (logged from training_batch.raw_latent_shape) batch_size = wandb_summary.get("batch_size") seq_len = wandb_summary.get("dit_seq_len") context_len = wandb_summary.get("context_len") @@ -213,7 +204,6 @@ def test_distributed_training(): world_size = int(NUM_GPUS_PER_NODE) total_peak_flops = peak_flops_per_gpu * world_size - # Calculate MFU achieved_flops_per_sec = achieved_flops / avg_step_time if avg_step_time > 0 else 0 mfu = (achieved_flops_per_sec / total_peak_flops * 100) if total_peak_flops > 0 else 0 diff --git a/fastvideo/tests/training/Vanilla/test_training_loss.py b/fastvideo/tests/training/Vanilla/test_training_loss.py index 6afe31169..6220e0386 100644 --- a/fastvideo/tests/training/Vanilla/test_training_loss.py +++ b/fastvideo/tests/training/Vanilla/test_training_loss.py @@ -10,7 +10,6 @@ from pathlib import Path import torch import json from huggingface_hub import snapshot_download -# Import the training pipeline sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) from fastvideo.training.runner import main @@ -27,10 +26,8 @@ NUM_GPUS_PER_NODE = "2" def run_worker(): """Worker function that will be run on each GPU""" - # Create and populate args parser = build_parser() - # Set the arguments as they are in finetune_v1_test.sh args = parser.parse_args([ "--pipeline_class", "WanTrainingPipeline", "--pipeline_module", "fastvideo.training.wan_training_pipeline", @@ -50,7 +47,6 @@ def run_worker(): "--not_apply_cfg_solver", "--dit_precision", "fp32", "--max_grad_norm", "1.0" ]) - # Call the main training function main(args) @@ -65,10 +61,8 @@ def test_distributed_training(): repo_type="dataset", local_dir_use_symlinks=False) - # Get the current file path current_file = Path(__file__).resolve() - # Run torchrun command cmd = [ "torchrun", "--nnodes", NUM_NODES, "--nproc_per_node", NUM_GPUS_PER_NODE, "--master_port", os.environ["MASTER_PORT"], diff --git a/fastvideo/tests/training/distill/test_distill_dmd.py b/fastvideo/tests/training/distill/test_distill_dmd.py index 54b68e5a0..d5f626f1d 100644 --- a/fastvideo/tests/training/distill/test_distill_dmd.py +++ b/fastvideo/tests/training/distill/test_distill_dmd.py @@ -10,7 +10,6 @@ from pathlib import Path import torch import json from huggingface_hub import snapshot_download -# Import the training pipeline sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) from fastvideo.training.runner import main @@ -24,10 +23,8 @@ NUM_GPUS_PER_NODE = "2" def run_worker(): """Worker function that will be run on each GPU""" - # Create and populate args parser = build_parser() - # Set the arguments as they are in finetune_v1_test.sh args = parser.parse_args([ "--pipeline_class", "WanDistillationPipeline", "--pipeline_module", "fastvideo.training.wan_distillation_pipeline", @@ -49,7 +46,6 @@ def run_worker(): "--enable_gradient_checkpointing_type", "full" ]) - # Call the main training function main(args) @@ -64,10 +60,8 @@ def test_distributed_training(): repo_type="dataset", local_dir_use_symlinks=False) - # Get the current file path current_file = Path(__file__).resolve() - # Run torchrun command cmd = [ "torchrun", "--nnodes", NUM_NODES, "--nproc_per_node", NUM_GPUS_PER_NODE, "--master_port", os.environ["MASTER_PORT"], diff --git a/fastvideo/tests/training/lora/test_lora_training.py b/fastvideo/tests/training/lora/test_lora_training.py index 67a5170e7..c097dc868 100644 --- a/fastvideo/tests/training/lora/test_lora_training.py +++ b/fastvideo/tests/training/lora/test_lora_training.py @@ -29,7 +29,6 @@ def test_lora_training(): repo_type="dataset", local_dir_use_symlinks=False) - # Run torchrun command directly on the training pipeline like the shell script cmd = [ "torchrun", "--nnodes", NUM_NODES, "--nproc_per_node", NUM_GPUS_PER_NODE, "--master_port", os.environ["MASTER_PORT"], "fastvideo/training/wan_training_pipeline.py", "--model_path", @@ -63,10 +62,8 @@ def test_lora_training(): wandb_summary = json.load(open(summary_file)) - # Define thresholds for LoRA training based on the provided console outputs fields_and_thresholds = { 'avg_step_time': 30.0, # something up with modal - # 'grad_norm': 0.05, # too volatile for now. TODO: fix nondeterminism in training 'step_time': 30.0, # something up with modal 'train_loss': 0.05 } diff --git a/fastvideo/tests/training/self-forcing/test_self_forcing.py b/fastvideo/tests/training/self-forcing/test_self_forcing.py index ac50d757f..df1f1bc8d 100644 --- a/fastvideo/tests/training/self-forcing/test_self_forcing.py +++ b/fastvideo/tests/training/self-forcing/test_self_forcing.py @@ -10,7 +10,6 @@ from pathlib import Path import torch import json from huggingface_hub import snapshot_download -# Import the training pipeline sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) from fastvideo.training.runner import main @@ -24,10 +23,8 @@ NUM_GPUS_PER_NODE = "2" def run_worker(): """Worker function that will be run on each GPU""" - # Create and populate args parser = build_parser() - # Set the arguments based on the distill_dmd_t2v_1.3B.sh script args = parser.parse_args([ "--pipeline_class", "WanSelfForcingDistillationPipeline", "--pipeline_module", "fastvideo.training.wan_self_forcing_distillation_pipeline", @@ -148,7 +145,6 @@ def run_worker(): "100", ]) - # Call the main training function main(args) @@ -163,10 +159,8 @@ def test_distributed_training(): repo_type="dataset", local_dir_use_symlinks=False) - # Get the current file path current_file = Path(__file__).resolve() - # Run torchrun command cmd = [ "torchrun", "--nnodes", NUM_NODES, "--nproc_per_node", NUM_GPUS_PER_NODE, "--master_port", os.environ["MASTER_PORT"], diff --git a/fastvideo/tests/transformers/test_cosmos.py b/fastvideo/tests/transformers/test_cosmos.py index eb2bdaae6..4a04d0ce5 100644 --- a/fastvideo/tests/transformers/test_cosmos.py +++ b/fastvideo/tests/transformers/test_cosmos.py @@ -76,11 +76,9 @@ def test_cosmos2_transformer(): weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2) logger.info("Weight mean difference: %s", weight_mean_diff) - # Set both models to eval mode model1 = model1.eval() model2 = model2.eval() - # Create identical inputs for both models batch_size = 1 seq_len = 30 @@ -107,7 +105,6 @@ def test_cosmos2_transformer(): # padding mask padding_mask = hidden_states.new_zeros(1, 1, 32, 32, device=device, dtype=precision) - # print(padding_mask.shape) forward_batch = ForwardBatch(data_type="dummy", ) @@ -129,7 +126,6 @@ def test_cosmos2_transformer(): timestep=timestep, padding_mask=padding_mask) - # Check if outputs have the same shape assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}" assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}" diff --git a/fastvideo/tests/transformers/test_cosmos2_5.py b/fastvideo/tests/transformers/test_cosmos2_5.py index 70161f91d..9ff950b7b 100644 --- a/fastvideo/tests/transformers/test_cosmos2_5.py +++ b/fastvideo/tests/transformers/test_cosmos2_5.py @@ -9,7 +9,6 @@ import sys import pytest import torch -# Add cosmos-predict2.5 to Python path for loading reference model TEST_DIR = os.path.dirname(os.path.abspath(__file__)) COSMOS_PREDICT2_5_PATH = os.path.join(TEST_DIR, '..', '..', '..', '..', 'cosmos-predict2.5') @@ -22,7 +21,6 @@ from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger from fastvideo.tests.utils import skip_if_gated_repo_inaccessible from fastvideo.utils import maybe_download_model -# Use Cosmos 2.5 specific config from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig from fastvideo.pipelines.pipeline_batch_info import ForwardBatch @@ -97,12 +95,10 @@ def load_reference_cosmos25_model(checkpoint_path: str, device, dtype): model = MinimalV1LVGDiT(**model_config) - # Load checkpoint if path exists if os.path.exists(checkpoint_path): logger.info(f"Loading reference model from {checkpoint_path}") checkpoint = torch.load(checkpoint_path, map_location='cpu') - # Extract state dict if 'state_dict' in checkpoint: checkpoint_state = checkpoint['state_dict'] elif 'model' in checkpoint: @@ -126,7 +122,6 @@ def load_reference_cosmos25_model(checkpoint_path: str, device, dtype): # Strip 'net.' prefix new_key = key[4:] if key.startswith('net.') else key # Strip 'net.' prefix - # Add '_checkpoint_wrapped_module' if needed if needs_checkpoint_wrapper and new_key.startswith('blocks.'): # Pattern: 'blocks.N.something' -> 'blocks.N._checkpoint_wrapped_module.something' parts = new_key.split('.', 2) @@ -135,7 +130,6 @@ def load_reference_cosmos25_model(checkpoint_path: str, device, dtype): transformed_state[new_key] = value - # Load with strict=False to handle any remaining mismatches missing_keys, unexpected_keys = model.load_state_dict(transformed_state, strict=False) if missing_keys: @@ -182,7 +176,6 @@ def test_cosmos25_transformer(): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") precision = torch.bfloat16 - # Create COSMOS 2.5 specific config from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25ArchConfig arch_config = Cosmos25ArchConfig( @@ -207,11 +200,9 @@ def test_cosmos25_transformer(): cosmos25_config = Cosmos25VideoConfig(arch_config=arch_config) - # Create FastVideo model directly (Cosmos 2.5 is not in diffusers format) logger.info("Creating FastVideo COSMOS 2.5 model...") from fastvideo.models.dits.cosmos2_5 import Cosmos25Transformer3DModel - # Get hf_config from the arch_config for model initialization hf_config = { 'in_channels': arch_config.in_channels, 'out_channels': arch_config.out_channels, @@ -236,7 +227,6 @@ def test_cosmos25_transformer(): fastvideo_model = fastvideo_model.to(device, dtype=precision) fastvideo_model.eval() - # Construct checkpoint path using relative paths checkpoint_file = os.path.join(MODEL_PATH, CHECKPOINT_SUBDIR, CHECKPOINT_FILENAME) if not os.path.exists(checkpoint_file): @@ -244,14 +234,12 @@ def test_cosmos25_transformer(): logger.info("Will test architecture without loading checkpoint weights") checkpoint_file = None - # Load checkpoint into FastVideo model using param_names_mapping if checkpoint_file: logger.info(f"Loading checkpoint into FastVideo model from {checkpoint_file}") from fastvideo.models.loader.utils import hf_to_custom_state_dict, get_param_names_mapping checkpoint = torch.load(checkpoint_file, map_location='cpu') - # Extract state dict (checkpoint might have 'state_dict', 'model', or be the dict itself) if 'state_dict' in checkpoint: checkpoint_state = checkpoint['state_dict'] elif 'model' in checkpoint: @@ -262,7 +250,6 @@ def test_cosmos25_transformer(): # Filter to only model parameters (remove training metadata like accum_*) model_state = {k: v for k, v in checkpoint_state.items() if k.startswith('net.') and 'accum_' not in k} - # Convert checkpoint keys to FastVideo format using param_names_mapping param_names_mapping_fn = get_param_names_mapping(cosmos25_config.arch_config.param_names_mapping) custom_state_dict, reverse_mapping = hf_to_custom_state_dict(model_state, param_names_mapping_fn) @@ -273,7 +260,6 @@ def test_cosmos25_transformer(): for k, v in custom_state_dict.items() if k in model_param_names } - # Load into FastVideo model missing_keys, unexpected_keys = fastvideo_model.load_state_dict(filtered_state_dict, strict=False) if missing_keys: @@ -292,12 +278,10 @@ def test_cosmos25_transformer(): logger.info("Loading reference COSMOS 2.5 model...") reference_model = load_reference_cosmos25_model(checkpoint_file, device, precision) if checkpoint_file else None - # Set models to eval mode fastvideo_model = fastvideo_model.eval() if reference_model is not None: reference_model = reference_model.eval() - # Create test inputs batch_size = 1 seq_len = 77 # Typical T5 sequence length @@ -359,7 +343,6 @@ def test_cosmos25_transformer(): # Reference model (if available) if reference_model is not None: - # Prepare input for reference model # MinimalV1LVGDiT adds condition mask internally, so pass them separately from cosmos_predict2._src.predict2.conditioner import DataType @@ -384,7 +367,6 @@ def test_cosmos25_transformer(): data_type=ref_data_type, ) - # Check FastVideo output shape and dtype logger.info(f"FastVideo output shape: {output_fv.shape}") logger.info(f"FastVideo output dtype: {output_fv.dtype}") assert output_fv.shape[0] == batch_size, "Batch size mismatch" @@ -395,13 +377,11 @@ def test_cosmos25_transformer(): if reference_model is not None: logger.info(f"Reference output shape: {output_ref.shape}") - # Check if outputs have the same shape assert output_fv.shape == output_ref.shape, \ f"Output shapes don't match: {output_fv.shape} vs {output_ref.shape}" assert output_fv.dtype == output_ref.dtype, \ f"Output dtype don't match: {output_fv.dtype} vs {output_ref.dtype}" - # Check if outputs are similar max_diff = torch.max(torch.abs(output_fv - output_ref)) mean_diff = torch.mean(torch.abs(output_fv - output_ref)) relative_diff = mean_diff / (torch.mean(torch.abs(output_ref)) + 1e-8) @@ -426,7 +406,6 @@ def test_cosmos25_transformer_video(): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") precision = torch.bfloat16 - # Create COSMOS 2.5 specific config from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25ArchConfig arch_config = Cosmos25ArchConfig( @@ -451,11 +430,9 @@ def test_cosmos25_transformer_video(): cosmos25_config = Cosmos25VideoConfig(arch_config=arch_config) - # Create FastVideo model directly (Cosmos 2.5 is not in diffusers format) logger.info("Creating FastVideo COSMOS 2.5 model for video test...") from fastvideo.models.dits.cosmos2_5 import Cosmos25Transformer3DModel - # Get hf_config from the arch_config for model initialization hf_config = { 'in_channels': arch_config.in_channels, 'out_channels': arch_config.out_channels, @@ -480,7 +457,6 @@ def test_cosmos25_transformer_video(): model = model.to(device, dtype=precision) model.eval() - # Create video input with multiple frames batch_size = 1 num_frames = 16 # Video with 16 frames seq_len = 77 @@ -495,7 +471,6 @@ def test_cosmos25_transformer_video(): dtype=precision) condition_mask = torch.zeros(batch_size, 1, num_frames, 64, 64, device=device, dtype=precision) - # Set first 2 frames as conditioning condition_mask[:, :, :2, :, :] = 1.0 encoder_hidden_states = torch.randn( @@ -534,7 +509,6 @@ def test_cosmos25_transformer_video(): logger.info(f"Video output shape: {output.shape}") logger.info(f"Video output dtype: {output.dtype}") - # Check output shape assert output.shape[0] == batch_size, "Batch size mismatch" assert output.shape[1] == 16, "Output channels should be 16" assert output.shape[2] == num_frames, "Number of frames mismatch" @@ -544,6 +518,5 @@ def test_cosmos25_transformer_video(): if __name__ == "__main__": - # Run tests directly test_cosmos25_transformer() test_cosmos25_transformer_video() diff --git a/fastvideo/tests/transformers/test_hyworld.py b/fastvideo/tests/transformers/test_hyworld.py index 321dea1bb..65a4f2b51 100644 --- a/fastvideo/tests/transformers/test_hyworld.py +++ b/fastvideo/tests/transformers/test_hyworld.py @@ -86,7 +86,6 @@ def test_hyworld_transformer(): torch.manual_seed(42) - # Create inputs for the model batch_size = 1 seq_len = 30 seq_len_2 = 70 @@ -103,7 +102,6 @@ def test_hyworld_transformer(): # Text embeddings [B, L, D] (including global token) encoder_hidden_states = torch.randn(batch_size, seq_len + 1, 3584, device=device, dtype=precision) - # Create attention mask for encoder_hidden_states encoder_attention_mask = torch.ones(batch_size, seq_len + 1, device=device, dtype=torch.bool) encoder_hidden_states[:, 15:] = 0 encoder_attention_mask[:, 15:] = False diff --git a/fastvideo/tests/utils.py b/fastvideo/tests/utils.py index d80c2198f..5e54022a6 100644 --- a/fastvideo/tests/utils.py +++ b/fastvideo/tests/utils.py @@ -116,7 +116,6 @@ def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True): frames1 = _read_visual_frames(video1_path) frames2 = _read_visual_frames(video2_path) - # Ensure same number of frames min_frames = min(frames1.shape[0], frames2.shape[0]) frames1 = frames1[:min_frames] frames2 = frames2[:min_frames] @@ -130,7 +129,6 @@ def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True): ssim_values = [] - # Process each frame individually for i in range(min_frames): img1 = frames1[i:i + 1] img2 = frames2[i:i + 1] diff --git a/fastvideo/tests/vaes/test_hunyuan15_vae.py b/fastvideo/tests/vaes/test_hunyuan15_vae.py index 707139426..0f18e39bf 100644 --- a/fastvideo/tests/vaes/test_hunyuan15_vae.py +++ b/fastvideo/tests/vaes/test_hunyuan15_vae.py @@ -10,8 +10,6 @@ from diffusers import AutoencoderKLHunyuanVideo15 import fastvideo.envs as envs from fastvideo.configs.pipelines import PipelineConfig from fastvideo.logger import init_logger -# from fastvideo.models.vaes.hunyuanvae import ( -# AutoencoderKLHunyuanVideo as MyHunyuanVAE) from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.models.loader.component_loader import VAELoader from fastvideo.configs.models.vaes import Hunyuan15VAEConfig diff --git a/fastvideo/tests/workflow/test_components_parquet_saver.py b/fastvideo/tests/workflow/test_components_parquet_saver.py index aceea9dbe..05c18968e 100644 --- a/fastvideo/tests/workflow/test_components_parquet_saver.py +++ b/fastvideo/tests/workflow/test_components_parquet_saver.py @@ -33,7 +33,6 @@ def test_parquet_dataset_saver_flush_and_last(tmp_path: Path): ]) B = 5 - # Build a minimal PreprocessBatch batch = PreprocessBatch( data_type="video", latents=torch.randn(B, 2), diff --git a/fastvideo/train/callbacks/validation.py b/fastvideo/train/callbacks/validation.py index a6c05603e..448fc8a0f 100644 --- a/fastvideo/train/callbacks/validation.py +++ b/fastvideo/train/callbacks/validation.py @@ -236,9 +236,6 @@ class ValidationCallback(Callback): return value.strip().lower() in {"1", "true", "yes", "on"} return bool(value) - # ---------------------------------------------------------- - # Callback hooks - # ---------------------------------------------------------- def _adopt_training_denoising_ladder( self, @@ -322,9 +319,6 @@ class ValidationCallback(Callback): self._run_validation(method, iteration) - # ---------------------------------------------------------- - # Core validation logic - # ---------------------------------------------------------- def _run_validation( self, @@ -845,9 +839,6 @@ class ValidationCallback(Callback): step, ) - # ---------------------------------------------------------- - # Metric evaluation - # ---------------------------------------------------------- def _metric_device(self) -> str: device = self.metrics_config.device @@ -1179,9 +1170,6 @@ class ValidationCallback(Callback): return -1 return int(self.metrics_config.mouse_pitch_sign) - # ---------------------------------------------------------- - # Pipeline management - # ---------------------------------------------------------- def _get_sampling_param(self) -> SamplingParam: if self._sampling_param is None: @@ -1341,9 +1329,6 @@ class ValidationCallback(Callback): self._pipeline_key = key return self._pipeline - # ---------------------------------------------------------- - # Batch preparation - # ---------------------------------------------------------- def _prepare_validation_batch( self, @@ -1465,9 +1450,6 @@ class ValidationCallback(Callback): tensor, ) - # ---------------------------------------------------------- - # Validation loop - # ---------------------------------------------------------- def _student_sampling_pipeline_config( self, @@ -1673,9 +1655,6 @@ class ValidationCallback(Callback): action=action, ) - # ---------------------------------------------------------- - # State management - # ---------------------------------------------------------- def state_dict(self) -> dict[str, Any]: state: dict[str, Any] = {} diff --git a/fastvideo/train/methods/base.py b/fastvideo/train/methods/base.py index 911414df8..49c1b4754 100644 --- a/fastvideo/train/methods/base.py +++ b/fastvideo/train/methods/base.py @@ -77,7 +77,6 @@ class TrainingMethod(torch.nn.Module, ABC): if mods: self.role_modules[role] = torch.nn.ModuleDict(mods) - # ------------------------------------------------------------------ def set_tracker(self, tracker: Any) -> None: self.tracker = tracker diff --git a/fastvideo/train/methods/consistency_model/causal_cd.py b/fastvideo/train/methods/consistency_model/causal_cd.py index b0c93640b..8f974f422 100644 --- a/fastvideo/train/methods/consistency_model/causal_cd.py +++ b/fastvideo/train/methods/consistency_model/causal_cd.py @@ -58,7 +58,6 @@ class CausalConsistencyDistillationMethod(TrainingMethod): ) self._init_optimizers_and_schedulers() - # ------------------------------------------------------------------ @property def _optimizer_dict(self) -> dict[str, Any]: @@ -85,7 +84,6 @@ class CausalConsistencyDistillationMethod(TrainingMethod): states["roles.ema.transformer"] = _FullModelState(self.ema_model.transformer) return states - # ------------------------------------------------------------------ def single_train_step( self, @@ -168,7 +166,6 @@ class CausalConsistencyDistillationMethod(TrainingMethod): metrics: dict[str, LogScalar] = {} return loss_map, outputs, metrics - # ------------------------------------------------------------------ def backward( self, @@ -189,9 +186,6 @@ class CausalConsistencyDistillationMethod(TrainingMethod): if iteration >= self._ema_start_step: self._update_ema() - # ------------------------------------------------------------------ - # Internal helpers - # ------------------------------------------------------------------ def _predict_flow( self, diff --git a/fastvideo/train/methods/distribution_matching/dmd2.py b/fastvideo/train/methods/distribution_matching/dmd2.py index ee1890bcb..8a771abb1 100644 --- a/fastvideo/train/methods/distribution_matching/dmd2.py +++ b/fastvideo/train/methods/distribution_matching/dmd2.py @@ -63,7 +63,6 @@ class DMD2Method(TrainingMethod): self._score_max_timestep, ) = self._parse_score_timestep_bounds() - # Initialize preprocessors on student. self.student.init_preprocessors(self.training_config) self._init_optimizers_and_schedulers() diff --git a/fastvideo/train/methods/distribution_matching/self_forcing.py b/fastvideo/train/methods/distribution_matching/self_forcing.py index 4939a85bd..b1e7f719b 100644 --- a/fastvideo/train/methods/distribution_matching/self_forcing.py +++ b/fastvideo/train/methods/distribution_matching/self_forcing.py @@ -55,7 +55,6 @@ class SelfForcingMethod(DMD2Method): role_models=role_models, ) - # Validate causal student. if not isinstance(self.student, CausalModelBase): raise ValueError("SelfForcingMethod requires a causal student " "implementing CausalModelBase.") diff --git a/fastvideo/train/methods/fine_tuning/dfsft.py b/fastvideo/train/methods/fine_tuning/dfsft.py index 3b08560fd..44dd78c87 100644 --- a/fastvideo/train/methods/fine_tuning/dfsft.py +++ b/fastvideo/train/methods/fine_tuning/dfsft.py @@ -37,7 +37,6 @@ class DiffusionForcingSFTMethod(TrainingMethod): self._timestep_index_range = (self._parse_timestep_index_range()) self._training_weights = self._build_training_weights() - # Initialize preprocessors on student. self.student.init_preprocessors(self.training_config) self._init_optimizers_and_schedulers() diff --git a/fastvideo/train/methods/fine_tuning/finetune.py b/fastvideo/train/methods/fine_tuning/finetune.py index c02c325a2..2b9f1d06a 100644 --- a/fastvideo/train/methods/fine_tuning/finetune.py +++ b/fastvideo/train/methods/fine_tuning/finetune.py @@ -104,7 +104,6 @@ class FineTuneMethod(TrainingMethod): "trainable") self._attn_kind: Literal["dense", "vsa"] = (self._infer_attn_kind()) - # Initialize preprocessors on student. self.student.init_preprocessors(self.training_config) self._init_optimizers_and_schedulers() diff --git a/fastvideo/train/methods/knowledge_distillation/kd.py b/fastvideo/train/methods/knowledge_distillation/kd.py index 45565fa9f..e7a1ecd97 100644 --- a/fastvideo/train/methods/knowledge_distillation/kd.py +++ b/fastvideo/train/methods/knowledge_distillation/kd.py @@ -61,9 +61,6 @@ from fastvideo.train.utils.optimizer import build_optimizer_and_scheduler logger = logging.getLogger(__name__) -# --------------------------------------------------------------------------- -# Cache helpers -# --------------------------------------------------------------------------- _COMPLETE_SENTINEL = "COMPLETE" _METADATA_FILE = "metadata.json" @@ -156,9 +153,6 @@ class _KDPathCache: dist.barrier() -# --------------------------------------------------------------------------- -# Dataset for reading cached paths at training time -# --------------------------------------------------------------------------- class _KDPathDataset(Dataset): @@ -187,9 +181,6 @@ class _KDPathDataset(Dataset): return torch.load(path, weights_only=True, map_location="cpu") -# --------------------------------------------------------------------------- -# Lazy dataloader wrapper -# --------------------------------------------------------------------------- class _KDDataLoaderWrapper: @@ -259,9 +250,6 @@ def _kd_collate(samples: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tenso return {k: torch.stack([s[k] for s in samples]) for k in keys} -# --------------------------------------------------------------------------- -# KD Method -# --------------------------------------------------------------------------- class KDMethod(TrainingMethod): @@ -332,7 +320,6 @@ class KDMethod(TrainingMethod): self._kd_wrapper = _KDDataLoaderWrapper(self._source_loader) self.student.dataloader = self._kd_wrapper # builder captures this - # --- Build SelfForcingFlowMatchScheduler for sigma lookups --- # num_inference_steps=1000 gives a dense grid accurate for any t. tc = self.training_config self._flow_shift = float(getattr(tc.pipeline_config, "flow_shift", 0.0) or 0.0) @@ -348,9 +335,6 @@ class KDMethod(TrainingMethod): # --- Student optimizer / scheduler (same as FineTuneMethod) --- self._init_optimizers_and_schedulers() - # ------------------------------------------------------------------ - # TrainingMethod abstract implementations - # ------------------------------------------------------------------ @property def _optimizer_dict(self) -> dict[str, Any]: @@ -368,9 +352,6 @@ class KDMethod(TrainingMethod): del iteration return [self._student_lr_scheduler] - # ------------------------------------------------------------------ - # Lifecycle - # ------------------------------------------------------------------ def on_train_start(self) -> None: super().on_train_start() # sets CUDA RNG, calls student.on_train_start() @@ -459,9 +440,6 @@ class KDMethod(TrainingMethod): for t in self._t_list[:-1]}, ) - # ------------------------------------------------------------------ - # Path generation - # ------------------------------------------------------------------ def _generate_paths( self, @@ -490,7 +468,6 @@ class KDMethod(TrainingMethod): if global_idx not in missing_set: continue # all ranks skip together — no FSDP divergence - # Load single sample by stable global index. raw = source_dataset.__getitems__([global_idx]) raw = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in raw.items()} @@ -547,7 +524,6 @@ class KDMethod(TrainingMethod): latents = training_batch.latents.to(dtype) B, T = latents.shape[:2] - # Build teacher's full ODE schedule. teacher_sched = SelfForcingFlowMatchScheduler( num_inference_steps=self._teacher_inference_steps, num_train_timesteps=int(self.student.num_train_timesteps), @@ -558,7 +534,6 @@ class KDMethod(TrainingMethod): ) teacher_t_list = [int(t.item()) for t in teacher_sched.timesteps] - # Start from pure noise at the highest teacher timestep. noise = torch.randn( latents.shape, device=device, @@ -627,9 +602,6 @@ class KDMethod(TrainingMethod): idx = torch.argmin((timesteps - t).abs()) return sigmas[idx].expand(numel).reshape(numel, 1, 1, 1) - # ------------------------------------------------------------------ - # Training step - # ------------------------------------------------------------------ def single_train_step( self, @@ -694,7 +666,6 @@ class KDMethod(TrainingMethod): attn_kind="dense", ) # [B, T, C, H, W] - # Convert predicted noise to predicted x0 using _sf_scheduler. timestep_bt = timestep_b.unsqueeze(1).expand(B, T) pred_x0 = pred_noise_to_pred_video( pred_noise=pred_noise.flatten(0, 1), @@ -740,9 +711,6 @@ class KDMethod(TrainingMethod): grad_accum_rounds=grad_accum_rounds, ) - # ------------------------------------------------------------------ - # Internal helpers - # ------------------------------------------------------------------ def _init_optimizers_and_schedulers(self) -> None: tc = self.training_config @@ -764,9 +732,6 @@ class KDMethod(TrainingMethod): ) -# --------------------------------------------------------------------------- -# Causal KD Method -# --------------------------------------------------------------------------- class KDCausalMethod(KDMethod): @@ -813,7 +778,6 @@ class KDCausalMethod(KDMethod): B, _S, T, C, H, W = batch["trajectory_latents"].shape K = self._num_steps # number of student steps (excludes t=0) - # Gather the K relevant trajectory steps for t_list[:-1]. traj_indices = torch.tensor( [self._t_to_traj_idx[t] for t in self._t_list[:-1]], dtype=torch.long, @@ -824,7 +788,6 @@ class KDCausalMethod(KDMethod): # Sample per-frame step indices, block-quantized. indexes = self._sample_per_frame_step_idx(B, T, K, device) # [B, T] in [0, K) - # Gather noisy_input per frame: [B, T, C, H, W] noisy_input = torch.gather( relevant, dim=1, diff --git a/fastvideo/train/models/cosmos/cosmos.py b/fastvideo/train/models/cosmos/cosmos.py index 2bb73d425..89cd4fca9 100644 --- a/fastvideo/train/models/cosmos/cosmos.py +++ b/fastvideo/train/models/cosmos/cosmos.py @@ -71,9 +71,6 @@ class CosmosModel(WanModel): transformer_override_safetensor=(transformer_override_safetensor), ) - # ------------------------------------------------------------------ - # Overrides - # ------------------------------------------------------------------ def prepare_batch( self, @@ -187,7 +184,6 @@ class CosmosModel(WanModel): # layout Cosmos's transformer expects. noisy_latents = noisy_latents.permute(0, 2, 1, 3, 4) - # ---------------------------------------------------------- # Cosmos 2.5: velocity prediction, no preconditioning. # # Noise schedule (EDM): x_t = x_0 + sigma * eps @@ -199,7 +195,6 @@ class CosmosModel(WanModel): # pred_x0 = x_t - pred * sigma # Since v ≈ noise, this gives: # pred_x0 = (x_0 + sigma*eps) - eps*sigma = x_0 ✓ - # ---------------------------------------------------------- assert batch.sigmas is not None sigma_1d = batch.sigmas.flatten()[:noisy_latents.shape[0]] diff --git a/fastvideo/train/models/hunyuan/hunyuan.py b/fastvideo/train/models/hunyuan/hunyuan.py index 2c2b17167..48ad4cada 100644 --- a/fastvideo/train/models/hunyuan/hunyuan.py +++ b/fastvideo/train/models/hunyuan/hunyuan.py @@ -65,9 +65,6 @@ class HunyuanModel(WanModel): lora=lora, ) - # ------------------------------------------------------------------ - # Overrides - # ------------------------------------------------------------------ def prepare_batch( self, @@ -197,7 +194,6 @@ class HunyuanModel(WanModel): model_path = maybe_download_model(tc.model_path) - # Use configured precisions for each encoder. precisions = tc.pipeline_config.text_encoder_precisions llama_dtype = PRECISION_TO_TYPE[precisions[0]] clip_dtype = PRECISION_TO_TYPE[precisions[1]] diff --git a/fastvideo/train/models/minimax_h3/minimax_h3_dmd.py b/fastvideo/train/models/minimax_h3/minimax_h3_dmd.py index 055369c5f..a37134858 100644 --- a/fastvideo/train/models/minimax_h3/minimax_h3_dmd.py +++ b/fastvideo/train/models/minimax_h3/minimax_h3_dmd.py @@ -199,9 +199,6 @@ class MiniMaxH3DMDModel(MiniMaxH3Model): def num_train_timesteps(self) -> int: return _DMD_TIMESTEP_SCALE - # ------------------------------------------------------------------ - # Packed dual-modality helpers - # ------------------------------------------------------------------ def _modality_shapes(self) -> tuple[tuple[int, int, int, int, int], tuple[int, int, int, int]]: """Return the ``[1, T, C, H, W]`` video and ``[1, 2, 32, Ta]`` audio shapes.""" @@ -363,9 +360,6 @@ class MiniMaxH3DMDModel(MiniMaxH3Model): ), ) - # ------------------------------------------------------------------ - # ModelBase overrides (packed convention) - # ------------------------------------------------------------------ def set_requires_negative_conditioning(self, requires: bool) -> None: """Fail fast: H3 cannot encode negative prompts at training time.""" @@ -731,9 +725,6 @@ class MiniMaxH3DMDModel(MiniMaxH3Model): sigma_fp64 = sigma.to(device=noisy.device, dtype=torch.float64) return (noisy_fp64 - sigma_fp64 * pred_noise_fp64).to(original_dtype) - # ------------------------------------------------------------------ - # Intermediate-latent visualization (LatentVisCallback) - # ------------------------------------------------------------------ def _load_vis_vae(self) -> Any: """Lazily load the H3 video VAE for visualization decodes. diff --git a/fastvideo/train/models/minimax_h3/minimax_h3_pdd.py b/fastvideo/train/models/minimax_h3/minimax_h3_pdd.py index d9f53f96c..034c11343 100644 --- a/fastvideo/train/models/minimax_h3/minimax_h3_pdd.py +++ b/fastvideo/train/models/minimax_h3/minimax_h3_pdd.py @@ -73,9 +73,6 @@ class MiniMaxH3PDDMixin(_PDDMixinHost): """Only this role widens its heads; the shared pipeline config stays single-head.""" return {"pdd_steps": self.pdd_steps} - # ------------------------------------------------------------------ - # PDD schedule contract - # ------------------------------------------------------------------ def pdd_fine_grid(self, device: torch.device) -> torch.Tensor: return pdd_fine_grid(self.pdd_steps, max_t=PDD_GRID_MAX_T, device=device) @@ -102,9 +99,6 @@ class MiniMaxH3PDDMixin(_PDDMixinHost): raise ValueError("MiniMaxH3PDDModel packs one document per batch; every sample must share a base time") return base * _DMD_TIMESTEP_SCALE - # ------------------------------------------------------------------ - # Forwards - # ------------------------------------------------------------------ def predict_pdd_heads( self, diff --git a/fastvideo/train/utils/moduleloader.py b/fastvideo/train/utils/moduleloader.py index 6c3449c2e..4475bc590 100644 --- a/fastvideo/train/utils/moduleloader.py +++ b/fastvideo/train/utils/moduleloader.py @@ -27,9 +27,7 @@ if TYPE_CHECKING: from fastvideo.train.utils.training_config import ( TrainingConfig, ) -# ------------------------------------------------------------------ # TrainingArgs builders (only place that creates FastVideoArgs) -# ------------------------------------------------------------------ def _make_training_args( @@ -85,9 +83,6 @@ def make_inference_args( return args -# ------------------------------------------------------------------ -# Module loading -# ------------------------------------------------------------------ def load_module_from_path( diff --git a/fastvideo/training/cosmos2_5_training_pipeline.py b/fastvideo/training/cosmos2_5_training_pipeline.py index 11fbe9d91..3fbde9f05 100644 --- a/fastvideo/training/cosmos2_5_training_pipeline.py +++ b/fastvideo/training/cosmos2_5_training_pipeline.py @@ -60,9 +60,6 @@ class Cosmos25TrainingPipeline(TrainingPipeline): ) self.validation_pipeline = validation_pipeline - # ------------------------------------------------------------------ - # Cosmos 2.5-specific overrides - # ------------------------------------------------------------------ def _normalize_dit_input(self, training_batch: TrainingBatch) -> TrainingBatch: """Skip: Cosmos25WanVAEWrapper normalises latents during encoding. @@ -129,9 +126,6 @@ class Cosmos25TrainingPipeline(TrainingPipeline): return training_batch -# --------------------------------------------------------------------------- -# Entry point (mirrors wan_training_pipeline.py) -# --------------------------------------------------------------------------- if __name__ == "__main__": logger.warning("\n" diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 524bffe79..d5f888b8e 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -97,7 +97,6 @@ class DistillationPipeline(TrainingPipeline): if training_args.real_score_model_path: logger.info("Loading real score transformer from: %s", training_args.real_score_model_path) training_args.override_transformer_cls_name = "WanTransformer3DModel" - # TODO(will): can use deepcopy instead if the model is the same self.real_score_transformer = self.load_module_from_path(training_args.real_score_model_path, "transformer", training_args) modules["real_score_transformer"] = self.real_score_transformer @@ -149,7 +148,6 @@ class DistillationPipeline(TrainingPipeline): else: self.boundary_timestep = None - # make sure the real score transformer is not trainable self.real_score_transformer.requires_grad_(False) self.real_score_transformer.eval() if self.real_score_transformer_2 is not None: @@ -169,10 +167,8 @@ class DistillationPipeline(TrainingPipeline): self.real_score_transformer_2 = apply_activation_checkpointing( self.real_score_transformer_2, checkpointing_type=training_args.enable_gradient_checkpointing_type) - # Initialize optimizers fake_score_params = list(filter(lambda p: p.requires_grad, self.fake_score_transformer.parameters())) - # Use separate learning rate for fake_score_transformer if specified fake_score_lr = training_args.fake_score_learning_rate if fake_score_lr == 0.0: fake_score_lr = training_args.learning_rate @@ -346,12 +342,10 @@ class DistillationPipeline(TrainingPipeline): return try: - # Save main transformer EMA if self.generator_ema is not None: ema_save_dir = os.path.join(output_dir, f"ema_checkpoint-{step}") os.makedirs(ema_save_dir, exist_ok=True) - # save as diffusers format from safetensors.torch import save_file from fastvideo.training.training_utils import (custom_to_hf_state_dict, gather_state_dict_on_cpu_rank0) @@ -376,12 +370,10 @@ class DistillationPipeline(TrainingPipeline): logger.info("EMA weights saved to %s", weight_path) - # Save transformer_2 EMA if self.generator_ema_2 is not None and self.transformer_2 is not None: ema_2_save_dir = os.path.join(output_dir, f"ema_2_checkpoint-{step}") os.makedirs(ema_2_save_dir, exist_ok=True) - # save as diffusers format from safetensors.torch import save_file from fastvideo.training.training_utils import (custom_to_hf_state_dict, gather_state_dict_on_cpu_rank0) @@ -528,13 +520,10 @@ class DistillationPipeline(TrainingPipeline): latents = training_batch.latents dtype = latents.dtype - # Step 1: Randomly sample a target timestep index from denoising_step_list target_timestep_idx = torch.randint(0, len(self.denoising_step_list), [1], device=self.device, dtype=torch.long) target_timestep_idx_int = target_timestep_idx.item() target_timestep = self.denoising_step_list[target_timestep_idx] - # Step 2: Simulate the multi-step inference process up to the target timestep - # Start from pure noise like in inference current_noise_latents = torch.randn(self.video_latent_shape, device=self.device, dtype=dtype) current_noise_latents_copy = current_noise_latents.clone() @@ -548,7 +537,6 @@ class DistillationPipeline(TrainingPipeline): for step_idx in range(max_target_idx): current_timestep = self.denoising_step_list[step_idx] current_timestep_tensor = current_timestep * torch.ones(1, device=self.device, dtype=torch.long) - # Run student model to get flow prediction training_batch_temp = self._build_distill_input_kwargs(current_noise_latents, current_timestep_tensor, training_batch.conditional_dict, @@ -560,7 +548,6 @@ class DistillationPipeline(TrainingPipeline): scheduler=self.noise_scheduler).unflatten( 0, pred_flow.shape[:2]) - # Add noise for the next timestep next_timestep = self.denoising_step_list[step_idx + 1] next_timestep_tensor = next_timestep * torch.ones(1, device=self.device, dtype=torch.long) noise = torch.randn(self.video_latent_shape, device=self.device, dtype=pred_clean.dtype) @@ -569,16 +556,12 @@ class DistillationPipeline(TrainingPipeline): latent_copy = current_noise_latents.clone() noise_latents.append(latent_copy) - # Step 3: Use the simulated noisy input for the final training step - # For timestep index 0, this is pure noise - # For timestep index k > 0, this is the result after k denoising steps + noise at target level if noise_latent_index >= 0: assert noise_latent_index < len(self.denoising_step_list) - 1, "noise_latent_index is out of bounds" noisy_input = noise_latents[noise_latent_index] else: noisy_input = current_noise_latents_copy - # Step 4: Final student prediction (this is what we train on) training_batch = self._build_distill_input_kwargs(noisy_input, target_timestep, training_batch.conditional_dict, training_batch) pred_noise = self.transformer(**training_batch.input_kwargs).permute(0, 2, 1, 3, 4) @@ -636,7 +619,6 @@ class DistillationPipeline(TrainingPipeline): # real_score_transformer uncond forward training_batch = self._build_distill_input_kwargs(noisy_latent, timestep, training_batch.unconditional_dict, training_batch) - # Use same transformer as conditional forward for consistency real_score_pred_noise_uncond = current_real_score_transformer(**training_batch.input_kwargs).permute( 0, 2, 1, 3, 4) @@ -762,9 +744,7 @@ class DistillationPipeline(TrainingPipeline): batch = next(self.train_loader_iter, None) # type: ignore if batch is None: self.current_epoch += 1 - # Reset iterator for next epoch self.train_loader_iter = iter(self.train_dataloader) - # Get first batch of new epoch batch = next(self.train_loader_iter) device = get_local_torch_device() @@ -972,7 +952,6 @@ class DistillationPipeline(TrainingPipeline): # First call parent class method to get basic training info super()._log_training_info() - # Then add distillation-specific information logger.info("Distillation-specific settings:") logger.info(" Generator update ratio: %s", self.generator_update_interval) assert isinstance(self.training_args, TrainingArgs) @@ -1017,14 +996,11 @@ class DistillationPipeline(TrainingPipeline): logger.info("Starting validation") - # Create sampling parameters if not provided sampling_param = SamplingParam.from_pretrained(training_args.model_path) - # Set deterministic seed for validation logger.info("Using validation seed: %s", self.seed) - # Prepare validation prompts logger.info('rank: %s: fastvideo_args.validation_dataset_file: %s', self.global_rank, training_args.validation_dataset_file, @@ -1032,7 +1008,6 @@ class DistillationPipeline(TrainingPipeline): validation_dataset = ValidationDataset(training_args.validation_dataset_file) validation_dataloader = DataLoader(validation_dataset, batch_size=None, num_workers=0) - # Set both transformers to eval mode transformer.eval() if hasattr(self, 'transformer_2') and self.transformer_2 is not None: self.transformer_2.eval() @@ -1049,7 +1024,6 @@ class DistillationPipeline(TrainingPipeline): if self.generator_ema is not None: ema_context = self.generator_ema.apply_to_model(validation_transformer) - # Handle transformer_2 EMA if available if hasattr(self, 'transformer_2') and self.transformer_2 is not None and self.generator_ema_2 is not None: ema_2_context = self.generator_ema_2.apply_to_model(self.transformer_2) logger.info("Using EMA_2 model for transformer_2 validation") @@ -1063,7 +1037,6 @@ class DistillationPipeline(TrainingPipeline): # Log validation results for this step world_group = get_world_group() num_sp_groups = world_group.world_size // self.sp_group.world_size - # Process each validation prompt for each validation step for num_inference_steps in validation_steps: logger.info("rank: %s: num_inference_steps: %s", self.global_rank, @@ -1102,14 +1075,12 @@ class DistillationPipeline(TrainingPipeline): assert batch.prompt is not None and isinstance(batch.prompt, str) captions.append(batch.prompt) - # Run validation inference with torch.no_grad(): output_batch = self.validation_pipeline.forward(batch, training_args) samples = output_batch.output.cpu() if self.rank_in_sp_group != 0: continue - # Process outputs video = rearrange(samples, "b c t h w -> t b c h w") frames = [] for x in video: @@ -1122,7 +1093,6 @@ class DistillationPipeline(TrainingPipeline): return videos, captions, audios, audio_sample_rates - # Apply EMA contexts if available (nested context managers) if ema_context is not None and ema_2_context is not None: with ema_context, ema_2_context: (step_videos, step_captions, step_audios, @@ -1231,7 +1201,6 @@ class DistillationPipeline(TrainingPipeline): latents = _prepare_vae_latents(latents) latents = _apply_vae_scale(latents) - # Apply shifting if needed if (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None): if isinstance(self.vae.shift_factor, torch.Tensor): latents += self.vae.shift_factor.to(latents.device, latents.dtype) @@ -1246,18 +1215,14 @@ class DistillationPipeline(TrainingPipeline): video_artifact = self.tracker.video(video, fps=24, format="mp4") # change to 16 for Wan2.1 if video_artifact is not None: tracker_loss_dict[latent_key] = video_artifact - # Clean up references del video, latents - # Process DMD training data if available - use decode_stage instead of self.vae.decode if 'generator_pred_video' in dmd_latents_vis_dict: for latent_key in dmd_log_keys: latents = dmd_latents_vis_dict[latent_key] latents = _prepare_vae_latents(latents) - # decoded_latent = decode_stage(ForwardBatch(data_type="video", latents=latents), training_args) latents = _apply_vae_scale(latents) - # Apply shifting if needed if (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None): if isinstance(self.vae.shift_factor, torch.Tensor): latents += self.vae.shift_factor.to(latents.device, latents.dtype) @@ -1272,7 +1237,6 @@ class DistillationPipeline(TrainingPipeline): video_artifact = self.tracker.video(video, fps=24, format="mp4") # change to 16 for Wan2.1 if video_artifact is not None: tracker_loss_dict[latent_key] = video_artifact - # Clean up references del video, latents # Log to tracker @@ -1286,21 +1250,17 @@ class DistillationPipeline(TrainingPipeline): # Set the same seed within each SP group to ensure reproducibility if self.sp_world_size > 1: - # Use the same seed for all processes within the same SP group sp_group_seed = seed + (self.global_rank // self.sp_world_size) set_random_seed(sp_group_seed) logger.info("Rank %s: Using SP group seed %s", self.global_rank, sp_group_seed) else: set_random_seed(seed + self.global_rank) - # Set random seeds for deterministic training self.noise_random_generator = torch.Generator(device="cpu").manual_seed(self.seed + self.global_rank) self.noise_gen_cuda = torch.Generator(device="cuda").manual_seed(self.seed + self.global_rank) self.validation_random_generator = torch.Generator(device="cpu").manual_seed(self.seed + self.global_rank) logger.info("Initialized random seeds with seed: %s", seed + self.global_rank) - # Initialize current_trainstep for EMA ready checks - #TODO: check if needed self.current_trainstep = self.init_steps # Resume from checkpoint if specified (this will restore random states) @@ -1373,7 +1333,6 @@ class DistillationPipeline(TrainingPipeline): progress_bar.update(1) if self.global_rank == 0: - # Prepare logging data log_data = { "train_total_loss": total_loss, "train_fake_score_loss": fake_score_loss, @@ -1414,7 +1373,6 @@ class DistillationPipeline(TrainingPipeline): self.tracker.log(log_data, step) - # Save training state checkpoint (for resuming training) if (self.training_args.training_state_checkpointing_steps > 0 and step % self.training_args.training_state_checkpointing_steps == 0): print("rank", self.global_rank, "save training state checkpoint at step", step) @@ -1477,7 +1435,6 @@ class DistillationPipeline(TrainingPipeline): self.tracker.finish() - # Save final training state checkpoint print("rank", self.global_rank, "save final training state checkpoint at step", self.training_args.max_train_steps) save_distillation_checkpoint( diff --git a/fastvideo/training/ltx2_training_pipeline.py b/fastvideo/training/ltx2_training_pipeline.py index c352bbff5..f7d9b9a9a 100644 --- a/fastvideo/training/ltx2_training_pipeline.py +++ b/fastvideo/training/ltx2_training_pipeline.py @@ -44,7 +44,6 @@ class LTX2TrainingPipeline(TrainingPipeline): _rope_fps: float = _DEFAULT_ROPE_FPS def initialize_pipeline(self, fastvideo_args: FastVideoArgs): - # TODO (David): Change to port LTX2 scheduler into self.modules["scheduler"] if "scheduler" in self.modules: del self.modules["scheduler"] diff --git a/fastvideo/training/matrixgame2_ode_causal_pipeline.py b/fastvideo/training/matrixgame2_ode_causal_pipeline.py index 88a885b8b..c172a2f63 100644 --- a/fastvideo/training/matrixgame2_ode_causal_pipeline.py +++ b/fastvideo/training/matrixgame2_ode_causal_pipeline.py @@ -145,7 +145,6 @@ class MatrixGame2ODEInitTrainingPipeline(TrainingPipeline): # [B, S, C, T, H, W] -> [B, S, T, C, H, W] to match self-forcing trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5) - # Move to device device = get_local_torch_device() training_batch.image_embeds = clip_feature.to(device, dtype=torch.bfloat16) training_batch.image_latents = first_frame_latent.to(device, dtype=torch.bfloat16) @@ -178,8 +177,6 @@ class MatrixGame2ODEInitTrainingPipeline(TrainingPipeline): max_timestep, [batch_size, num_frame], device=self.device, dtype=torch.long) - # logger.info(f"individual timestep: {timestep}") - # make the noise level the same within every block timestep = timestep.reshape(timestep.shape[0], -1, num_frame_per_block) timestep[:, :, 1:] = timestep[:, :, 0:1] timestep = timestep.reshape(timestep.shape[0], -1) @@ -242,7 +239,6 @@ class MatrixGame2ODEInitTrainingPipeline(TrainingPipeline): dim=1) timestep = self.dmd_denoising_steps[indexes] - # Prepare inputs for transformer latent_vis_dict["noisy_input"] = noisy_input.permute(0, 2, 1, 3, 4).detach().clone().cpu() latent_vis_dict["x0"] = target_latent.permute(0, 2, 1, 3, 4).detach().clone().cpu() @@ -300,7 +296,6 @@ class MatrixGame2ODEInitTrainingPipeline(TrainingPipeline): mask = t != 0 - # Compute loss loss = F.mse_loss(noise_pred[mask], target_latent[mask], reduction="mean") loss = loss / args.gradient_accumulation_steps @@ -387,7 +382,6 @@ class MatrixGame2ODEInitTrainingPipeline(TrainingPipeline): video_artifact = self.tracker.video(video, fps=16, format="mp4") # change to 16 for Wan2.1 if video_artifact is not None: tracker_loss_dict[latent_key] = video_artifact - # Clean up references del video, pixel_latent, latent if self.global_rank == 0 and tracker_loss_dict: diff --git a/fastvideo/training/matrixgame2_self_forcing_distillation_pipeline.py b/fastvideo/training/matrixgame2_self_forcing_distillation_pipeline.py index 7dcb55cd7..24b48fa18 100644 --- a/fastvideo/training/matrixgame2_self_forcing_distillation_pipeline.py +++ b/fastvideo/training/matrixgame2_self_forcing_distillation_pipeline.py @@ -166,7 +166,6 @@ class MatrixGame2SelfForcingDistillationPipeline(SelfForcingDistillationPipeline frame_seq_length = post_patch_height * post_patch_width self.frame_seq_length = frame_seq_length - # Get model configuration parameters - handle FSDP wrapping num_attention_heads = getattr(self.transformer, 'num_attention_heads', None) attention_head_dim = getattr(self.transformer, 'attention_head_dim', None) @@ -199,7 +198,6 @@ class MatrixGame2SelfForcingDistillationPipeline(SelfForcingDistillationPipeline torch.tensor([0], dtype=torch.long, device=device) }) - # Initialize cross-attention cache crossattn_cache = [] for _ in range(num_transformer_blocks): crossattn_cache.append({ @@ -213,7 +211,6 @@ class MatrixGame2SelfForcingDistillationPipeline(SelfForcingDistillationPipeline False }) - # Initialize action module KV caches action_heads_num = action_config.get('heads_num', 16) if action_config else 16 mouse_hidden_dim = action_config.get('mouse_hidden_dim', 1024) if action_config else 1024 keyboard_hidden_dim = action_config.get('keyboard_hidden_dim', 1024) if action_config else 1024 @@ -357,11 +354,9 @@ class MatrixGame2SelfForcingDistillationPipeline(SelfForcingDistillationPipeline device=noise.device, dtype=noise.dtype) - # Step 1: Initialize KV cache to all zeros (self.kv_cache1, self.crossattn_cache, self.kv_cache_mouse, self.kv_cache_keyboard) = self._initialize_simulation_caches(batch_size, dtype, self.device) - # Step 2: Temporal denoising loop current_start_frame = 0 all_num_frames = [self.num_frame_per_block] * num_blocks if self.independent_first_frame: @@ -373,7 +368,6 @@ class MatrixGame2SelfForcingDistillationPipeline(SelfForcingDistillationPipeline for block_index, current_num_frames in enumerate(all_num_frames): noisy_input = noise[:, current_start_frame:current_start_frame + current_num_frames] - # Step 3.1: Spatial denoising loop for index, current_timestep in enumerate(self.denoising_step_list): exit_flag = (index == exit_flags[0]) if self.same_step_across_blocks else (index == exit_flags[block_index]) @@ -388,7 +382,6 @@ class MatrixGame2SelfForcingDistillationPipeline(SelfForcingDistillationPipeline if not exit_flag: with torch.no_grad(): - # Build input kwargs training_batch_temp = self._build_distill_input_kwargs(noisy_input, timestep, training_batch.conditional_dict, @@ -468,10 +461,8 @@ class MatrixGame2SelfForcingDistillationPipeline(SelfForcingDistillationPipeline 0, pred_flow.shape[:2]) break - # Step 3.2: record the model's output output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred - # Step 3.3: rerun with timestep zero to update the cache context_timestep = torch.ones_like(timestep) * self.context_noise denoised_pred = self.noise_scheduler.add_noise(denoised_pred.flatten(0, 1), torch.randn_like(denoised_pred.flatten(0, 1)), @@ -497,10 +488,8 @@ class MatrixGame2SelfForcingDistillationPipeline(SelfForcingDistillationPipeline current_start=current_start_frame * self.frame_seq_length, start_frame=current_start_frame) - # Step 3.4: update the start and end frame indices current_start_frame += current_num_frames - # Handle last 21 frames logic pred_image_or_video = output # Slice last 21 frames if we generated more @@ -512,7 +501,6 @@ class MatrixGame2SelfForcingDistillationPipeline(SelfForcingDistillationPipeline # Decode to video latent_to_decode = latent_to_decode.permute(0, 2, 1, 3, 4) # [B, C, F, H, W] - # Apply VAE scaling and shift factors if isinstance(self.vae.scaling_factor, torch.Tensor): latent_to_decode = latent_to_decode / self.vae.scaling_factor.to( latent_to_decode.device, latent_to_decode.dtype) @@ -537,7 +525,6 @@ class MatrixGame2SelfForcingDistillationPipeline(SelfForcingDistillationPipeline else: pred_image_or_video_last_21 = pred_image_or_video - # Set up gradient mask if we generated more than minimum frames if num_generated_frames != min_num_frames: # Currently, we do not use gradient for the first chunk, since it contains image latents gradient_mask = torch.ones_like(pred_image_or_video_last_21, dtype=torch.bool) @@ -546,24 +533,20 @@ class MatrixGame2SelfForcingDistillationPipeline(SelfForcingDistillationPipeline else: gradient_mask[:, :self.num_frame_per_block] = False - # Apply gradient masking if needed final_output = pred_image_or_video_last_21.to(dtype) if gradient_mask is not None: - # Apply gradient masking: detach frames that shouldn't contribute gradients final_output = torch.where( gradient_mask, pred_image_or_video_last_21, # Keep original values where gradient_mask is True pred_image_or_video_last_21.detach() # Detach where gradient_mask is False ) - # Store visualization data training_batch.dmd_latent_vis_dict["generator_timestep"] = torch.as_tensor( self.denoising_step_list[exit_flags[0]], dtype=torch.float32, device=self.device, ).detach().clone() - # Store gradient mask information for debugging if gradient_mask is not None: training_batch.dmd_latent_vis_dict["gradient_mask"] = gradient_mask.float() training_batch.dmd_latent_vis_dict["num_generated_frames"] = torch.tensor(num_generated_frames, @@ -573,7 +556,6 @@ class MatrixGame2SelfForcingDistillationPipeline(SelfForcingDistillationPipeline dtype=torch.float32, device=self.device) - # Clean up caches assert self.kv_cache1 is not None assert self.crossattn_cache is not None self._reset_simulation_caches(self.kv_cache1, self.crossattn_cache, self.kv_cache_mouse, self.kv_cache_keyboard) @@ -620,9 +602,7 @@ class MatrixGame2SelfForcingDistillationPipeline(SelfForcingDistillationPipeline batch = next(self.train_loader_iter, None) # type: ignore if batch is None: self.current_epoch += 1 - # Reset iterator for next epoch self.train_loader_iter = iter(self.train_dataloader) - # Get first batch of new epoch batch = next(self.train_loader_iter) clip_feature = batch['clip_feature'] diff --git a/fastvideo/training/matrixgame2_training_pipeline.py b/fastvideo/training/matrixgame2_training_pipeline.py index 4ba256ae6..8f882f45e 100644 --- a/fastvideo/training/matrixgame2_training_pipeline.py +++ b/fastvideo/training/matrixgame2_training_pipeline.py @@ -47,8 +47,6 @@ class MatrixGame2TrainingPipeline(TrainingPipeline): args_copy.inference_mode = True args_copy.dit_cpu_offload = True - # args_copy.pipeline_config.vae_config.load_encoder = False - # validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained( self.validation_pipeline = MatrixGame2I2VPipeline.from_pretrained(training_args.model_path, args=None, inference_mode=True, @@ -66,15 +64,11 @@ class MatrixGame2TrainingPipeline(TrainingPipeline): if batch is None: self.current_epoch += 1 logger.info("Starting epoch %s", self.current_epoch) - # Reset iterator for next epoch self.train_loader_iter = iter(self.train_dataloader) - # Get first batch of new epoch batch = next(self.train_loader_iter) latents = batch['vae_latent'] latents = latents[:, :, :self.training_args.num_latent_t] - # encoder_hidden_states = batch['text_embedding'] - # encoder_attention_mask = batch['text_attention_mask'] clip_features = batch['clip_feature'] image_latents = batch['first_frame_latent'] image_latents = image_latents[:, :, :self.training_args.num_latent_t] diff --git a/fastvideo/training/ode_causal_pipeline.py b/fastvideo/training/ode_causal_pipeline.py index 84b5ecc65..e9cb5dd83 100644 --- a/fastvideo/training/ode_causal_pipeline.py +++ b/fastvideo/training/ode_causal_pipeline.py @@ -134,7 +134,6 @@ class ODEInitTrainingPipeline(TrainingPipeline): # [B, S, C, T, H, W] -> [B, S, T, C, H, W] to match self-forcing trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5) - # Move to device device = get_local_torch_device() training_batch.encoder_hidden_states = encoder_hidden_states.to(device, dtype=torch.bfloat16) training_batch.encoder_attention_mask = encoder_attention_mask.to(device, dtype=torch.bfloat16) @@ -159,8 +158,6 @@ class ODEInitTrainingPipeline(TrainingPipeline): max_timestep, [batch_size, num_frame], device=self.device, dtype=torch.long) - # logger.info(f"individual timestep: {timestep}") - # make the noise level the same within every block timestep = timestep.reshape(timestep.shape[0], -1, num_frame_per_block) timestep[:, :, 1:] = timestep[:, :, 0:1] timestep = timestep.reshape(timestep.shape[0], -1) @@ -192,13 +189,11 @@ class ODEInitTrainingPipeline(TrainingPipeline): dim=1, index=self._cached_closest_idx_per_dmd.to(traj_latents.device)) logger.info("relevant_traj_latents: %s", relevant_traj_latents.shape) - # assert relevant_traj_latents.shape[0] == 1 indexes = self._get_timestep( # [B, num_frames] 0, len(self.dmd_denoising_steps), B, num_frames, 3, uniform_timestep=False) logger.info("indexes: %s", indexes.shape) logger.info("indexes: %s", indexes) - # noisy_input = relevant_traj_latents[indexes] noisy_input = torch.gather(relevant_traj_latents, dim=1, index=indexes.reshape(B, 1, num_frames, 1, 1, @@ -207,7 +202,6 @@ class ODEInitTrainingPipeline(TrainingPipeline): timestep = self.dmd_denoising_steps[indexes] logger.info("selected timestep for rank %s: %s", self.global_rank, timestep, local_main_process_only=False) - # Prepare inputs for transformer latent_vis_dict["noisy_input"] = noisy_input.permute(0, 2, 1, 3, 4).detach().clone().cpu() latent_vis_dict["x0"] = target_latent.permute(0, 2, 1, 3, 4).detach().clone().cpu() @@ -257,7 +251,6 @@ class ODEInitTrainingPipeline(TrainingPipeline): mask = t != 0 - # Compute loss loss = F.mse_loss(noise_pred[mask], target_latent[mask], reduction="mean") loss = loss / args.gradient_accumulation_steps @@ -302,7 +295,6 @@ class ODEInitTrainingPipeline(TrainingPipeline): video_artifact = self.tracker.video(video, fps=16, format="mp4") # change to 16 for Wan2.1 if video_artifact is not None: tracker_loss_dict[latent_key] = video_artifact - # Clean up references del video, pixel_latent, latent if self.global_rank == 0 and tracker_loss_dict: diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index e1db17c70..8ab5bc3d0 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -93,7 +93,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): rank = dist.get_rank() if dist.is_initialized() else 0 if rank == 0: - # Generate random indices indices = torch.randint(low=0, high=num_denoising_steps, size=(num_blocks, ), device=device) if self.last_step_only: indices = torch.ones_like(indices) * (num_denoising_steps - 1) @@ -172,7 +171,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): num_generated_frames += 1 min_num_frames += 1 - # Create noise with dynamic shape if initial_latent is not None: noise_shape = [batch_size, num_generated_frames - 1, *self.video_latent_shape[2:]] else: @@ -207,20 +205,17 @@ class SelfForcingDistillationPipeline(DistillationPipeline): except (StopIteration, AttributeError): return "Unknown" - # Step 1: Initialize KV cache to all zeros cache_frames = num_generated_frames + num_input_frames self.kv_cache1, self.crossattn_cache = self._initialize_simulation_caches(batch_size, dtype, self.device, max_num_frames=cache_frames) - # Step 2: Cache context feature current_start_frame = 0 if initial_latent is not None: timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0 output[:, :1] = initial_latent with torch.no_grad(): - # Build input kwargs for initial latent training_batch_temp = self._build_distill_input_kwargs(initial_latent, timestep * 0, training_batch.conditional_dict, training_batch) # we process the image latent with self.transformer_2 (low-noise expert) @@ -236,7 +231,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): start_frame=current_start_frame) current_start_frame += 1 - # Step 3: Temporal denoising loop all_num_frames = [self.num_frame_per_block] * num_blocks if self.independent_first_frame and initial_latent is None: all_num_frames = [1] + all_num_frames @@ -248,7 +242,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): noisy_input = noise[:, current_start_frame - num_input_frames:current_start_frame + current_num_frames - num_input_frames] - # Step 3.1: Spatial denoising loop for index, current_timestep in enumerate(self.denoising_step_list): exit_flag = index == exit_flags[0] if self.same_step_across_blocks else index == exit_flags[block_index] @@ -262,7 +255,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): if not exit_flag: with torch.no_grad(): - # Build input kwargs training_batch_temp = self._build_distill_input_kwargs(noisy_input, timestep, training_batch.conditional_dict, training_batch) @@ -331,10 +323,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline): 0, pred_flow.shape[:2]) break - # Step 3.2: record the model's output output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred - # Step 3.3: rerun with timestep zero to update the cache context_timestep = torch.ones_like(timestep) * self.context_noise denoised_pred = self.noise_scheduler.add_noise(denoised_pred.flatten(0, 1), torch.randn_like(denoised_pred.flatten(0, 1)), @@ -356,10 +346,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline): current_start=current_start_frame * self.frame_seq_length, start_frame=current_start_frame) - # Step 3.4: update the start and end frame indices current_start_frame += current_num_frames - # Handle last 21 frames logic pred_image_or_video = output if num_input_frames > 0: pred_image_or_video = output[:, num_input_frames:] @@ -373,7 +361,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # Decode to video latent_to_decode = latent_to_decode.permute(0, 2, 1, 3, 4) # [B, C, F, H, W] - # Apply VAE scaling and shift factors if isinstance(self.vae.scaling_factor, torch.Tensor): latent_to_decode = latent_to_decode / self.vae.scaling_factor.to( latent_to_decode.device, latent_to_decode.dtype) @@ -398,7 +385,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): else: pred_image_or_video_last_21 = pred_image_or_video - # Set up gradient mask if we generated more than minimum frames if num_generated_frames != min_num_frames: # Currently, we do not use gradient for the first chunk, since it contains image latents gradient_mask = torch.ones_like(pred_image_or_video_last_21, dtype=torch.bool) @@ -407,22 +393,18 @@ class SelfForcingDistillationPipeline(DistillationPipeline): else: gradient_mask[:, :self.num_frame_per_block] = False - # Apply gradient masking if needed final_output = pred_image_or_video_last_21.to(dtype) if gradient_mask is not None: - # Apply gradient masking: detach frames that shouldn't contribute gradients final_output = torch.where( gradient_mask, pred_image_or_video_last_21, # Keep original values where gradient_mask is True pred_image_or_video_last_21.detach() # Detach where gradient_mask is False ) - # Store visualization data training_batch.dmd_latent_vis_dict["generator_timestep"] = torch.tensor(self.denoising_step_list[exit_flags[0]], dtype=torch.float32, device=self.device) - # Store gradient mask information for debugging if gradient_mask is not None: training_batch.dmd_latent_vis_dict["gradient_mask"] = gradient_mask.float() training_batch.dmd_latent_vis_dict["num_generated_frames"] = torch.tensor(num_generated_frames, @@ -432,7 +414,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): dtype=torch.float32, device=self.device) - # Clean up caches assert self.kv_cache1 is not None assert self.crossattn_cache is not None self._reset_simulation_caches(self.kv_cache1, self.crossattn_cache) @@ -459,7 +440,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): frame_seq_length = post_patch_height * post_patch_width self.frame_seq_length = frame_seq_length - # Get model configuration parameters - handle FSDP wrapping num_attention_heads = getattr(self.transformer, 'num_attention_heads', None) attention_head_dim = getattr(self.transformer, 'attention_head_dim', None) text_len = getattr(self.transformer, 'text_len', None) @@ -486,7 +466,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): torch.tensor([0], dtype=torch.long, device=device) }) - # Initialize cross-attention cache crossattn_cache = [] for _ in range(num_transformer_blocks): crossattn_cache.append({ @@ -527,9 +506,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): batch = next(self.train_loader_iter, None) # type: ignore if batch is None: self.current_epoch += 1 - # Reset iterator for next epoch self.train_loader_iter = iter(self.train_dataloader) - # Get first batch of new epoch batch = next(self.train_loader_iter) # latents, encoder_hidden_states, encoder_attention_mask, infos = batch @@ -585,7 +562,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): generator_log_dict = {} for batch in batches: - # Create a new batch with detached tensors batch_gen = TrainingBatch() for key, value in batch.__dict__.items(): if isinstance(value, torch.Tensor): @@ -604,7 +580,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): (generator_loss / gradient_accumulation_steps).backward() total_generator_loss += generator_loss.detach().item() generator_log_dict.update(gen_log_dict) - # Store visualization data from generator training if hasattr(batch_gen, 'dmd_latent_vis_dict'): training_batch.dmd_latent_vis_dict.update(batch_gen.dmd_latent_vis_dict) @@ -639,7 +614,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): critic_log_dict = {} for batch in batches: - # Create a new batch with detached tensors batch_critic = TrainingBatch() for key, value in batch.__dict__.items(): if isinstance(value, torch.Tensor): @@ -658,7 +632,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): (critic_loss / gradient_accumulation_steps).backward() total_critic_loss += critic_loss.detach().item() critic_log_dict.update(crit_log_dict) - # Store visualization data from critic training if hasattr(batch_critic, 'fake_score_latent_vis_dict'): training_batch.fake_score_latent_vis_dict.update(batch_critic.fake_score_latent_vis_dict) @@ -695,7 +668,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): if hasattr(training_batch, 'fake_score_latent_vis_dict'): logger.info("Fake score latent keys: %s", list(training_batch.fake_score_latent_vis_dict.keys())) - # Process generator predictions if available if hasattr(training_batch, 'dmd_latent_vis_dict') and training_batch.dmd_latent_vis_dict: dmd_latents_vis_dict = training_batch.dmd_latent_vis_dict dmd_log_keys = ['generator_pred_video', 'real_score_pred_video', 'faker_score_pred_video'] @@ -733,7 +705,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): tracker_loss_dict[f"dmd_{latent_key}"] = video_artifact del video, latents - # Process critic predictions if hasattr(training_batch, 'fake_score_latent_vis_dict') and training_batch.fake_score_latent_vis_dict: fake_score_latents_vis_dict = training_batch.fake_score_latent_vis_dict fake_score_log_keys = ['generator_pred_video'] @@ -799,7 +770,6 @@ class SelfForcingDistillationPipeline(DistillationPipeline): # Set the same seed within each SP group to ensure reproducibility if self.sp_world_size > 1: - # Use the same seed for all processes within the same SP group sp_group_seed = seed + (self.global_rank // self.sp_world_size) set_random_seed(sp_group_seed) else: diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index f5a3939ad..9a139cce6 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -105,7 +105,6 @@ class TrainingPipeline(LoRAPipeline, ABC): self.seed = training_args.seed self.set_schemas() - # Set random seeds for deterministic training assert self.seed is not None, "seed must be set" set_random_seed(self.seed + self.global_rank) self.transformer.train() @@ -120,7 +119,6 @@ class TrainingPipeline(LoRAPipeline, ABC): self.set_trainable() params_to_optimize = self.transformer.parameters() params_to_optimize = list(filter(lambda p: p.requires_grad, params_to_optimize)) - # Parse betas from string format "beta1,beta2" betas_str = training_args.betas betas = tuple(float(x.strip()) for x in betas_str.split(",")) @@ -194,7 +192,6 @@ class TrainingPipeline(LoRAPipeline, ABC): training_args.train_sp_batch_size) self.num_train_epochs = math.ceil(training_args.max_train_steps / self.num_update_steps_per_epoch) - # TODO(will): is there a cleaner way to track epochs? self.current_epoch = 0 trackers = list(training_args.trackers) @@ -235,9 +232,7 @@ class TrainingPipeline(LoRAPipeline, ABC): if batch is None: self.current_epoch += 1 logger.info("Starting epoch %s", self.current_epoch) - # Reset iterator for next epoch self.train_loader_iter = iter(self.train_dataloader) - # Get first batch of new epoch batch = next(self.train_loader_iter) latents = batch['vae_latent'] @@ -266,7 +261,6 @@ class TrainingPipeline(LoRAPipeline, ABC): return training_batch def _normalize_dit_input(self, training_batch: TrainingBatch) -> TrainingBatch: - # TODO(will): support other models with self.tracker.timed("timing/normalize_input"): training_batch.latents = normalize_dit_input( 'wan', @@ -334,10 +328,6 @@ class TrainingPipeline(LoRAPipeline, ABC): boundary_ratio = self.training_args.boundary_ratio if self.train_transformer_2: u = (1 - boundary_ratio) + u * boundary_ratio # min: 1 - boundary_ratio, max: 1 - # elif self.transformer_2 is not None: - # u = u * (1 - boundary_ratio) # min: 0, max: 1 - boundary_ratio - # else: # patch for now to align with non-MoE timestep logic - # pass indices = (u * self.noise_scheduler.config.num_train_timesteps).long() return self.noise_scheduler.timesteps[indices].to(device=device) @@ -396,11 +386,6 @@ class TrainingPipeline(LoRAPipeline, ABC): assert training_batch.attn_metadata is None input_kwargs = training_batch.input_kwargs - # if 'hunyuan' in self.training_args.model_type: - # input_kwargs["guidance"] = torch.tensor( - # [1000.0], - # device=training_batch.noisy_model_input.device, - # dtype=torch.bfloat16) current_model = self.transformer_2 if self.train_transformer_2 else self.transformer with self.tracker.timed("timing/forward_backward"), set_forward_context( @@ -436,9 +421,6 @@ class TrainingPipeline(LoRAPipeline, ABC): def _clip_grad_norm(self, training_batch: TrainingBatch) -> TrainingBatch: max_grad_norm = self.training_args.max_grad_norm - # TODO(will): perhaps move this into transformer api so that we can do - # the following: - # grad_norm = transformer.clip_grad_norm_(max_grad_norm) if max_grad_norm is not None: with self.tracker.timed("timing/clip_grad_norm"): # Only clip gradients for the model that is currently training @@ -469,9 +451,7 @@ class TrainingPipeline(LoRAPipeline, ABC): for _ in range(self.training_args.gradient_accumulation_steps): training_batch = self._get_next_batch(training_batch) - # Normalize DIT input training_batch = self._normalize_dit_input(training_batch) - # Create noisy model input training_batch = self._prepare_dit_inputs(training_batch) # old sharding code, need to shard latents and noise but not input @@ -528,7 +508,6 @@ class TrainingPipeline(LoRAPipeline, ABC): logger.info("Transformer 2: Starting training with %s B trainable parameters", round(num_trainable_params / 1e9, 3)) - # Set random seeds for deterministic training self.noise_random_generator = torch.Generator(device="cpu").manual_seed(self.seed + self.global_rank) self.noise_gen_cuda = torch.Generator(device=current_platform.device_name).manual_seed(self.seed + self.global_rank) @@ -564,7 +543,6 @@ class TrainingPipeline(LoRAPipeline, ABC): current_decay_times = min(step // vsa_decay_interval_steps, vsa_sparsity // vsa_decay_rate) current_vsa_sparsity = current_decay_times * vsa_decay_rate elif vmoba_available: - #TODO: add vmoba sparsity scheduling here current_vsa_sparsity = 0.0 else: current_vsa_sparsity = 0.0 @@ -714,10 +692,8 @@ class TrainingPipeline(LoRAPipeline, ABC): logger.info("Starting validation") - # Create sampling parameters if not provided sampling_param = SamplingParam.from_pretrained(training_args.model_path) - # Prepare validation prompts logger.info('rank: %s: fastvideo_args.validation_dataset_file: %s', self.global_rank, training_args.validation_dataset_file, @@ -736,7 +712,6 @@ class TrainingPipeline(LoRAPipeline, ABC): world_group = get_world_group() num_sp_groups = world_group.world_size // self.sp_group.world_size - # Process each validation prompt for each validation step for num_inference_steps in validation_steps: logger.info("rank: %s: num_inference_steps: %s", self.global_rank, @@ -762,7 +737,6 @@ class TrainingPipeline(LoRAPipeline, ABC): step_captions.append(batch.prompt) step_ref_videos.append(validation_batch.get("ref_video")) - # Run validation inference output_batch = self.validation_pipeline.forward(batch, training_args) samples = output_batch.output.cpu() @@ -779,7 +753,6 @@ class TrainingPipeline(LoRAPipeline, ABC): if self.rank_in_sp_group != 0: continue - # Process outputs video = rearrange(samples, "b c t h w -> t b c h w") frames = [] for x in video: @@ -904,21 +877,17 @@ class TrainingPipeline(LoRAPipeline, ABC): out_path = os.path.join(tmpdir, "muxed.mp4") wav_path = os.path.join(tmpdir, "audio.wav") - # Write audio to WAV file with wave.open(wav_path, "wb") as wav_file: wav_file.setnchannels(num_channels) wav_file.setsampwidth(2) wav_file.setframerate(sample_rate) wav_file.writeframes(audio_int16.tobytes()) - # Open input video and audio input_video = av.open(video_path) input_audio = av.open(wav_path) - # Create output with both streams output = av.open(out_path, mode="w") - # Add video stream (copy codec from input) in_video_stream = input_video.streams.video[0] out_video_stream = output.add_stream( codec_name=in_video_stream.codec_context.name, @@ -928,7 +897,6 @@ class TrainingPipeline(LoRAPipeline, ABC): out_video_stream.height = in_video_stream.height out_video_stream.pix_fmt = in_video_stream.pix_fmt - # Add audio stream (AAC) out_audio_stream = output.add_stream("aac", rate=sample_rate) out_audio_stream.layout = layout diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index 52d84744d..169a204e9 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -47,7 +47,6 @@ def gather_state_dict_on_cpu_rank0( if hasattr(param, "_local_tensor"): # DTensor case if param.is_cpu: - # Gather directly on CPU param = param.full_tensor() else: if device is not None: @@ -155,24 +154,20 @@ def save_checkpoint(transformer, cpu_state = gather_state_dict_on_cpu_rank0(transformer, device=None) if rank == 0: - # Save model weights (consolidated) transformer_save_dir = os.path.join(save_dir, "transformer") os.makedirs(transformer_save_dir, exist_ok=True) weight_path = os.path.join(transformer_save_dir, "diffusion_pytorch_model.safetensors") logger.info("rank: %s, saving consolidated checkpoint to %s", rank, weight_path, local_main_process_only=False) - # Convert training format to diffusers format and save diffusers_state_dict = custom_to_hf_state_dict(cpu_state, transformer.reverse_param_names_mapping) save_file(diffusers_state_dict, weight_path) logger.info("rank: %s, consolidated checkpoint saved to %s", rank, weight_path, local_main_process_only=False) - # Save model config config_dict = transformer.hf_config if "dtype" in config_dict: del config_dict["dtype"] # TODO config_path = os.path.join(transformer_save_dir, "config.json") - # save dict as json with open(config_path, "w") as f: json.dump(config_dict, f, indent=4) logger.info("--> checkpoint saved at step %s to %s", step, weight_path) @@ -224,12 +219,10 @@ def save_distillation_checkpoint( save_dir = os.path.join(output_dir, f"checkpoint-{step}") os.makedirs(save_dir, exist_ok=True) - # Create directories for models inference_save_dir = os.path.join(save_dir, "generator_inference_transformer") # Only save distributed checkpoint if not only saving generator weight if not only_save_generator_weight: - # Save generator distributed checkpoint generator_states = { "model": ModelWrapper(generator_transformer), } @@ -255,7 +248,6 @@ def save_distillation_checkpoint( end_time - begin_time, local_main_process_only=False) - # Save generator_2 distributed checkpoint (MoE support) if generator_transformer_2 is not None: generator_2_states = { "model": ModelWrapper(generator_transformer_2), @@ -282,7 +274,6 @@ def save_distillation_checkpoint( end_time - begin_time, local_main_process_only=False) - # Save critic distributed checkpoint critic_states = { "model": ModelWrapper(fake_score_transformer), } @@ -308,7 +299,6 @@ def save_distillation_checkpoint( end_time - begin_time, local_main_process_only=False) - # Save critic_2 distributed checkpoint (MoE support) if fake_score_transformer_2 is not None: critic_2_states = { "model": ModelWrapper(fake_score_transformer_2), @@ -335,7 +325,6 @@ def save_distillation_checkpoint( end_time - begin_time, local_main_process_only=False) - # Save real_score_transformer_2 distributed checkpoint (MoE support) if real_score_transformer_2 is not None: real_score_2_states = { "model": ModelWrapper(real_score_transformer_2), @@ -360,7 +349,6 @@ def save_distillation_checkpoint( end_time - begin_time, local_main_process_only=False) - # Save shared random state separately shared_states = { "random_state": RandomStateWrapper(noise_generator), } @@ -381,7 +369,6 @@ def save_distillation_checkpoint( if generator_ema is not None and getattr(generator_ema, "mode", None) == "rank0_full": _save_rank0_full_ema_safetensors(generator_ema, generator_transformer, rank, save_dir, "generator_ema") elif generator_ema is not None and getattr(generator_ema, "mode", None) == "local_shard": - # Save per-rank shard ema_dir_shard = os.path.join(save_dir, "ema_local_shard") os.makedirs(ema_dir_shard, exist_ok=True) ema_shard_path = os.path.join(ema_dir_shard, f"generator_ema_rank{rank}.pt") @@ -403,7 +390,6 @@ def save_distillation_checkpoint( _save_rank0_full_ema_safetensors(generator_ema_2, generator_transformer_2, rank, save_dir, "generator_ema_2") elif generator_ema_2 is not None and getattr(generator_ema_2, "mode", None) == "local_shard": - # Save per-rank shard for EMA_2 ema_dir_shard_2 = os.path.join(save_dir, "ema_local_shard") os.makedirs(ema_dir_shard_2, exist_ok=True) ema2_shard_path = os.path.join(ema_dir_shard_2, f"generator_ema_2_rank{rank}.pt") @@ -419,11 +405,9 @@ def save_distillation_checkpoint( except Exception as e: logger.warning("rank: %s, failed saving EMA_2 separately: %s", rank, str(e)) - # Save generator model weights (consolidated) for inference cpu_state = gather_state_dict_on_cpu_rank0(generator_transformer, device=None) if rank == 0: - # Save generator model weights (consolidated) for inference os.makedirs(inference_save_dir, exist_ok=True) weight_path = os.path.join(inference_save_dir, "diffusion_pytorch_model.safetensors") logger.info("rank: %s, saving consolidated generator inference checkpoint to %s", @@ -431,7 +415,6 @@ def save_distillation_checkpoint( weight_path, local_main_process_only=False) - # Convert training format to diffusers format and save diffusers_state_dict = custom_to_hf_state_dict(cpu_state, generator_transformer.reverse_param_names_mapping) save_file(diffusers_state_dict, weight_path) @@ -440,17 +423,14 @@ def save_distillation_checkpoint( weight_path, local_main_process_only=False) - # Save model config config_dict = generator_transformer.hf_config if "dtype" in config_dict: del config_dict["dtype"] # TODO config_path = os.path.join(inference_save_dir, "config.json") - # save dict as json with open(config_path, "w") as f: json.dump(config_dict, f, indent=4) logger.info("--> distillation checkpoint saved at step %s to %s", step, weight_path) - # Save generator_2 model weights (consolidated) for inference (MoE support) if generator_transformer_2 is not None: inference_save_dir_2 = os.path.join(save_dir, "generator_2_inference_transformer") cpu_state_2 = gather_state_dict_on_cpu_rank0(generator_transformer_2, device=None) @@ -463,7 +443,6 @@ def save_distillation_checkpoint( weight_path_2, local_main_process_only=False) - # Convert training format to diffusers format and save diffusers_state_dict_2 = custom_to_hf_state_dict(cpu_state_2, generator_transformer_2.reverse_param_names_mapping) save_file(diffusers_state_dict_2, weight_path_2) @@ -473,7 +452,6 @@ def save_distillation_checkpoint( weight_path_2, local_main_process_only=False) - # Save model config config_dict_2 = generator_transformer_2.hf_config if "dtype" in config_dict_2: del config_dict_2["dtype"] # TODO @@ -498,7 +476,6 @@ def load_checkpoint(transformer, logger.warning("Checkpoint path %s does not exist", checkpoint_path) return 0 - # Extract step number from checkpoint path step = int(os.path.basename(os.path.normpath(checkpoint_path)).split('-')[-1]) if rank == 0: @@ -581,13 +558,11 @@ def load_distillation_checkpoint( logger.warning("Distillation checkpoint path %s does not exist", checkpoint_path) return 0 - # Extract step number from checkpoint path step = int(os.path.basename(checkpoint_path).split('-')[-1]) if rank == 0: logger.info("Loading distillation checkpoint from step %s", step) - # Load generator distributed checkpoint generator_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint", "generator") if not os.path.exists(generator_dcp_dir): logger.warning("Generator distributed checkpoint directory %s does not exist", generator_dcp_dir) @@ -620,7 +595,6 @@ def load_distillation_checkpoint( end_time - begin_time, local_main_process_only=False) - # Load EMA separately if generator_ema is not None: try: if getattr(generator_ema, "mode", None) == "local_shard": @@ -637,7 +611,6 @@ def load_distillation_checkpoint( except Exception as e: logger.warning("rank: %s, failed to load generator EMA: %s", rank, str(e)) - # Load EMA_2 from its shard (symmetric with generator_ema above) if generator_ema_2 is not None: try: if getattr(generator_ema_2, "mode", None) == "local_shard": @@ -653,7 +626,6 @@ def load_distillation_checkpoint( except Exception as e: logger.warning("rank: %s, failed to load generator EMA_2: %s", rank, str(e)) - # Load generator_2 distributed checkpoint (MoE support) if generator_transformer_2 is not None: generator_2_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint", "generator_2") if os.path.exists(generator_2_dcp_dir): @@ -686,7 +658,6 @@ def load_distillation_checkpoint( else: logger.info("rank: %s, generator_2 checkpoint not found, skipping", rank) - # Load critic distributed checkpoint critic_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint", "critic") if not os.path.exists(critic_dcp_dir): logger.warning("Critic distributed checkpoint directory %s does not exist", critic_dcp_dir) @@ -719,7 +690,6 @@ def load_distillation_checkpoint( end_time - begin_time, local_main_process_only=False) - # Load critic_2 distributed checkpoint (MoE support) if fake_score_transformer_2 is not None: critic_2_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint", "critic_2") if os.path.exists(critic_2_dcp_dir): @@ -752,7 +722,6 @@ def load_distillation_checkpoint( else: logger.info("rank: %s, critic_2 checkpoint not found, skipping", rank) - # Load real_score_2 distributed checkpoint (MoE support) if real_score_transformer_2 is not None: real_score_2_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint", "real_score_2") if os.path.exists(real_score_2_dcp_dir): @@ -779,7 +748,6 @@ def load_distillation_checkpoint( else: logger.info("rank: %s, real_score_2 checkpoint not found, skipping", rank) - # Load shared random state shared_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint", "shared") if not os.path.exists(shared_dcp_dir): logger.warning("Shared random state directory %s does not exist", shared_dcp_dir) @@ -906,8 +874,6 @@ def clip_grad_norm_( """ grads = [p.grad for p in parameters if p.grad is not None] - # TODO(aryan): Wait for next Pytorch release to use `torch.nn.utils.get_total_norm` - # total_norm = torch.nn.utils.get_total_norm(grads, norm_type, error_if_nonfinite, foreach) total_norm = _get_total_norm(grads, norm_type, error_if_nonfinite, foreach) # If total_norm is a DTensor, the placements must be `torch.distributed._tensor.ops.math_ops._NormPartial`. @@ -1404,7 +1370,6 @@ def get_cosine_schedule_with_min_lr(optimizer: Optimizer, # Cosine decay from 1.0 to min_lr_ratio over num_cycles periods # Use the same formula as standard cosine but ensure minimum is min_lr_ratio instead of 0 cosine_value = 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)) - # Ensure the value doesn't go below min_lr_ratio return max(min_lr_ratio, cosine_value) return LambdaLR(optimizer, lr_lambda, last_epoch) @@ -1741,7 +1706,6 @@ class EMA_FSDP: for name, p in self.module.named_parameters(): if not p.requires_grad: continue - # Save local shard p_local = EMA_FSDP._to_local_tensor(p.detach()) if p_local.numel() == 0: # Nothing to swap on this rank for this param @@ -1752,7 +1716,6 @@ class EMA_FSDP: if ema_cpu.numel() != p_local.numel(): # Shard shape mismatch (e.g., empty shard here), skip continue - # Copy EMA shard into local param shard p_local.copy_(ema_cpu.to(dtype=p_local.dtype, device=p_local.device)) return self.module diff --git a/fastvideo/training/wan_i2v_distillation_pipeline.py b/fastvideo/training/wan_i2v_distillation_pipeline.py index 3161322ae..8a931752f 100644 --- a/fastvideo/training/wan_i2v_distillation_pipeline.py +++ b/fastvideo/training/wan_i2v_distillation_pipeline.py @@ -49,8 +49,6 @@ class WanI2VDistillationPipeline(DistillationPipeline): args_copy.inference_mode = True args_copy.dit_cpu_offload = False - # args_copy.pipeline_config.vae_config.load_encoder = False - # validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained( validation_pipeline = WanImageToVideoDmdPipeline.from_pretrained( training_args.model_path, args=None, @@ -68,9 +66,7 @@ class WanI2VDistillationPipeline(DistillationPipeline): if batch is None: self.current_epoch += 1 logger.info("Starting epoch %s", self.current_epoch) - # Reset iterator for next epoch self.train_loader_iter = iter(self.train_dataloader) - # Get first batch of new epoch batch = next(self.train_loader_iter) latents = batch['vae_latent'] diff --git a/fastvideo/training/wan_i2v_training_pipeline.py b/fastvideo/training/wan_i2v_training_pipeline.py index 8a8153858..50f70cbd2 100644 --- a/fastvideo/training/wan_i2v_training_pipeline.py +++ b/fastvideo/training/wan_i2v_training_pipeline.py @@ -47,8 +47,6 @@ class WanI2VTrainingPipeline(TrainingPipeline): args_copy.inference_mode = True args_copy.dit_cpu_offload = True - # args_copy.pipeline_config.vae_config.load_encoder = False - # validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained( self.validation_pipeline = WanImageToVideoPipeline.from_pretrained(training_args.model_path, args=None, inference_mode=True, @@ -66,9 +64,7 @@ class WanI2VTrainingPipeline(TrainingPipeline): if batch is None: self.current_epoch += 1 logger.info("Starting epoch %s", self.current_epoch) - # Reset iterator for next epoch self.train_loader_iter = iter(self.train_dataloader) - # Get first batch of new epoch batch = next(self.train_loader_iter) latents = batch['vae_latent'] diff --git a/fastvideo/utils.py b/fastvideo/utils.py index 3b7d29c92..c4240d57f 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -46,8 +46,6 @@ logger = init_logger(__name__) T = TypeVar("T") -# TODO(will): used to convert fastvideo_args.precision to torch.dtype. Find a -# cleaner way to do this. PRECISION_TO_TYPE = { "fp32": torch.float32, "fp16": torch.float16, @@ -184,7 +182,6 @@ class FlexibleArgumentParser(argparse.ArgumentParser): _DEFER_CONFIG_SUBCOMMANDS = frozenset({"generate", "serve"}) def __init__(self, *args, **kwargs) -> None: - # Set the default 'formatter_class' to SortedHelpFormatter if 'formatter_class' not in kwargs: kwargs['formatter_class'] = SortedHelpFormatter super().__init__(*args, **kwargs) @@ -204,7 +201,6 @@ class FlexibleArgumentParser(argparse.ArgumentParser): if '--config' in args and not self._should_defer_config_loading(args): args = self._pull_args_from_config(args) - # Convert underscores to dashes and vice versa in argument names processed_args = [] for arg in args: if arg.startswith('--'): @@ -236,12 +232,10 @@ class FlexibleArgumentParser(argparse.ArgumentParser): while i < len(args): arg = args[i] if arg.startswith('--'): - # Handle --key=value format if '=' in arg: key = arg.split('=')[0][2:].replace('-', '_') namespace._provided.add(key) i += 1 - # Handle --key value format else: key = arg[2:].replace('-', '_') namespace._provided.add(key) @@ -419,7 +413,6 @@ def warn_for_unimplemented_methods(cls: type[T]) -> type[T]: try: attr = getattr(self, attr_name) - # get the func of callable method if callable(attr): attr_func = attr.__func__ except AttributeError: @@ -731,18 +724,15 @@ def maybe_download_model_index(model_name_or_path: str, revision: str | None = N filename = f"{subfolder}/{config_filename}" if subfolder else config_filename model_index_path = hf_hub_download(repo_id=repo_id, filename=filename, revision=revision) - # Load the selected manifest. with open(model_index_path) as f: config: dict[str, Any] = json.load(f) - # Verify it has the required fields if "_class_name" not in config: raise ValueError(f"{config_filename} for {model_name_or_path} does not contain _class_name field") if "_diffusers_version" not in config: raise ValueError(f"{config_filename} for {model_name_or_path} does not contain _diffusers_version field") - # Add the pipeline name for downstream use config["pipeline_name"] = config["_class_name"] logger.info("Downloaded %s for %s, pipeline: %s", config_filename, model_name_or_path, config["_class_name"]) @@ -840,18 +830,12 @@ def shallow_asdict(obj) -> dict[str, Any]: return {f.name: getattr(obj, f.name) for f in fields(obj)} -# TODO: validate that this is fine def kill_itself_when_parent_died() -> None: - # if sys.platform == "linux": - # sigkill this process when parent worker manager dies PR_SET_PDEATHSIG = 1 import platform if platform.system() == "Linux": libc = ctypes.CDLL("libc.so.6") libc.prctl(PR_SET_PDEATHSIG, signal.SIGKILL) - # elif platform.system() == "Darwin": - # libc = ctypes.CDLL("libc.dylib") - # logger.warning("kill_itself_when_parent_died is only supported in linux.") else: logger.warning("kill_itself_when_parent_died is only supported in linux.") @@ -957,7 +941,6 @@ def dict_to_3d_list( "If mask_strategy is None, you must provide t_max, l_max, and h_max") return [[[None for _ in range(h_max)] for _ in range(l_max)] for _ in range(t_max)] - # Parse all keys into integer tuples indices = [tuple(map(int, key.split("_"))) for key in mask_strategy] # Decide on dimensions @@ -983,7 +966,6 @@ def dict_to_3d_list( t, l, h = map(int, key.split("_")) # noqa: E741 if 0 <= t < max_timesteps_idx and 0 <= l < max_layer_idx and 0 <= h < max_head_idx: result[t][l][h] = value - # else: silently ignore any key that doesn't fit return result @@ -1043,13 +1025,11 @@ def best_output_size(w, h, dw, dh, expected_area): ow = (expected_area * ratio)**0.5 oh = expected_area / ow - # process width first ow1 = int(ow // dw * dw) oh1 = int(expected_area / ow1 // dh * dh) assert ow1 % dw == 0 and oh1 % dh == 0 and ow1 * oh1 <= expected_area ratio1 = ow1 / oh1 - # process height first oh2 = int(oh // dh * dh) ow2 = int(expected_area / oh2 // dw * dw) assert oh2 % dh == 0 and ow2 % dw == 0 and ow2 * oh2 <= expected_area @@ -1077,7 +1057,6 @@ def pixels_to_uint8(pixels: torch.Tensor) -> torch.Tensor: def save_decoded_latents_as_video(decoded_latents: list[torch.Tensor], output_path: str, fps: int): - # Process outputs videos = rearrange(decoded_latents, "b c t h w -> t b c h w") frames = [] for x in videos: @@ -1156,7 +1135,6 @@ def _append_to_memory_trace(message: str, log_file_path: str | os.PathLike[str] pass -# TODO(xingyu): add adopted message for this def get_ip() -> str: host_ip = envs.FASTVIDEO_HOST_IP.get() if host_ip: diff --git a/fastvideo/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py index 7a80d8a9e..0045a511a 100644 --- a/fastvideo/worker/gpu_worker.py +++ b/fastvideo/worker/gpu_worker.py @@ -69,16 +69,6 @@ class Worker: self.rank = rank self.distributed_init_method = distributed_init_method - # Init request dispatcher - # TODO(will): add request dispatcher: use TypeBasedDispatcher from - # utils.py - # self._request_dispatcher = TypeBasedDispatcher( - # [ - # (RpcReqInput, self.handle_rpc_request), - # (GenerateRequest, self.handle_generate_request), - # (ExpertDistributionReq, self.expert_distribution_handle), - # ] - # ) def init_device(self) -> None: """Initialize the device for the worker.""" @@ -138,7 +128,6 @@ class Worker: device_id = self.device.index if self.device.index is not None else 0 self.fastvideo_args.finalize_device_offload_policy(device_id) - # Initialize the distributed environment. dist_timeout = (timedelta( seconds=self.fastvideo_args.dist_timeout) if self.fastvideo_args.dist_timeout is not None else None) sp_size = self.fastvideo_args.sp_size @@ -194,9 +183,7 @@ class Worker: def shutdown(self) -> dict[str, Any]: """Gracefully shut down the worker process""" logger.info("Worker %d shutting down...", self.rank, local_main_process_only=False) - # Clean up resources if hasattr(self, 'pipeline') and self.pipeline is not None: - # Clean up pipeline resources if needed pass # Destroy the distributed environment diff --git a/fastvideo/worker/multiproc_executor.py b/fastvideo/worker/multiproc_executor.py index fd3e51609..0e02a6768 100644 --- a/fastvideo/worker/multiproc_executor.py +++ b/fastvideo/worker/multiproc_executor.py @@ -148,7 +148,6 @@ class MultiprocExecutor(Executor): set_multiproc_executor_envs() - # Check if master_port is provided in fastvideo_args master_port = get_open_port(self.fastvideo_args.master_port) distributed_init_method = get_distributed_init_method(get_loopback_ip(), master_port) logger.info("Use master port: %s", master_port) @@ -188,7 +187,6 @@ class MultiprocExecutor(Executor): # Close death_writers first to signal workers to exit self._ensure_worker_termination([uw.proc for uw in unready_workers]) - # Register shutdown on exit atexit.register(self.shutdown) def execute_forward(self, forward_batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: @@ -202,7 +200,6 @@ class MultiprocExecutor(Executor): logging_info = None logging_info = responses[0]["logging_info"] if envs.FASTVIDEO_STAGE_LOGGING.get() else None - # Get extra dict (contains audio, peak_memory_mb, etc.) extra = responses[0].get("extra", {}) result_batch = ForwardBatch(data_type=forward_batch.data_type, @@ -399,7 +396,6 @@ class MultiprocExecutor(Executor): logger.info("Shutting down MultiprocExecutor...") - # Check if workers were initialized (they might not be if initialization failed) if not hasattr(self, 'workers') or not self.workers: logger.info("No workers to shut down.") return @@ -450,7 +446,6 @@ class MultiprocExecutor(Executor): if worker.proc.is_alive(): worker.proc.kill() - # Clean up pipes for worker in self.workers: with contextlib.suppress(Exception): worker.pipe.close() @@ -560,10 +555,8 @@ class WorkerMultiprocProc: wrapper.init_worker(all_kwargs) self.worker = wrapper - # Initialize device self.worker.init_device() - # Set process title and log prefix self.setup_proc_title_and_log_prefix() @staticmethod @@ -591,7 +584,6 @@ class WorkerMultiprocProc: "streaming_output_queue": streaming_output_queue, "log_queue": log_queue, } - # Run EngineCore busy loop in background process. proc = context.Process(target=WorkerMultiprocProc.worker_main, kwargs=process_kwargs, name=f"FVWorkerProc-{rank}", @@ -707,7 +699,6 @@ class WorkerMultiprocProc: finally: if ready_pipe is not None: ready_pipe.close() - # Clean up once worker exits busy loop if worker is not None: worker.shutdown() @@ -755,7 +746,6 @@ class WorkerMultiprocProc: raise e from None finally: - # Close connection. pipe.close() # If any workers failed, raise exception with details diff --git a/fastvideo/worker/ray_distributed_executor.py b/fastvideo/worker/ray_distributed_executor.py index bbed0b5b4..60e3f068f 100644 --- a/fastvideo/worker/ray_distributed_executor.py +++ b/fastvideo/worker/ray_distributed_executor.py @@ -170,7 +170,6 @@ class RayDistributedExecutor(Executor): # The remaining workers are the actual ray actors. self.workers: list[RayWorkerWrapper] = [] - # Create the workers. # use the first N bundles that have GPU resources. bundle_indices: list[int] = [] for bundle_id, bundle in enumerate(placement_group.bundle_specs): @@ -249,7 +248,6 @@ class RayDistributedExecutor(Executor): # Stamped before init_worker pickles the args to every worker. self._setup_h3_encoder_split([item.ip for item in sorted_worker_metadata]) - # Get the set of GPU IDs used on each node. worker_node_and_gpu_ids = self._run_ray_workers("get_node_and_gpu_ids") node_workers = defaultdict(list) # node id -> list of worker ranks node_gpus = defaultdict(list) # node id -> list of gpu ids @@ -278,7 +276,6 @@ class RayDistributedExecutor(Executor): " environment variable, make sure it is unique for" " each node.") - # Set environment variables for the driver and workers. all_args_to_update_environment_variables: list[dict[str, str]] if keep_raylet_devices: # The inherited device list stays; each worker gets its ordinal in it. @@ -298,9 +295,7 @@ class RayDistributedExecutor(Executor): destination="workers", ) - # Copy existing env vars to each worker's args for args in all_args_to_update_environment_variables: - # TODO: refactor platform-specific env vars for name in env_vars_to_copy: if name in os.environ: args[name] = os.environ[name] @@ -314,7 +309,6 @@ class RayDistributedExecutor(Executor): driver_ip = get_loopback_ip() distributed_init_method = get_distributed_init_method(driver_ip, get_open_port()) - # Initialize the actual workers inside worker wrapper. all_kwargs = [] for rank, (node_id, _) in enumerate(worker_node_and_gpu_ids): local_rank = node_workers[node_id].index(rank) @@ -507,11 +501,9 @@ class RayDistributedExecutor(Executor): sent_method = method if isinstance(method, str) else cloudpickle.dumps(method) del method - # Start the ray workers first. ray_workers = self.workers ray_worker_outputs = [worker.execute_method.remote(sent_method, *args, **kwargs) for worker in ray_workers] - # Get the results of the ray workers. ray_worker_outputs = ray.get(ray_worker_outputs) return ray_worker_outputs diff --git a/fastvideo/worker/ray_utils.py b/fastvideo/worker/ray_utils.py index a2f696e78..afb1fb0d7 100644 --- a/fastvideo/worker/ray_utils.py +++ b/fastvideo/worker/ray_utils.py @@ -180,7 +180,6 @@ def initialize_ray_cluster( raise ValueError(f"current platform {current_platform.device_name} does not " "support ray.") - # Create or get the placement group for worker processes current_placement_group = fastvideo_args.ray_placement_group or ray.util.get_current_placement_group() if current_placement_group: @@ -188,7 +187,6 @@ def initialize_ray_cluster( # We are in a placement group bundles = current_placement_group.bundle_specs - # Verify that we can use the placement group. device_bundles = 0 for bundle in bundles: bundle_devices = bundle.get(device_str, 0) @@ -213,7 +211,6 @@ def initialize_ray_cluster( logger.warning( "The number of required %ss exceeds the total " "number of available %ss in the placement group.", device_str, device_str) - # Create a new placement group placement_group_specs: list[dict[str, float]] = ([{device_str: 1.0} for _ in range(fastvideo_args.num_gpus)]) # FastVideo engine is also a worker to execute model with an accelerator, @@ -237,7 +234,6 @@ def initialize_ray_cluster( assert current_placement_group is not None _verify_bundles(current_placement_group, fastvideo_args, device_str) - # Set the placement group in the fastvideo args fastvideo_args.ray_placement_group = current_placement_group diff --git a/fastvideo/workflow/preprocess/components.py b/fastvideo/workflow/preprocess/components.py index ca9f1d4d0..47986ffed 100644 --- a/fastvideo/workflow/preprocess/components.py +++ b/fastvideo/workflow/preprocess/components.py @@ -155,14 +155,12 @@ class ParquetDatasetSaver: assert isinstance(batch.prompt_embeds, list) assert isinstance(batch.prompt_attention_mask, list) - # Process non-padded embeddings (if needed) if batch.prompt_attention_mask is not None: batch.prompt_embeds = self._process_non_padded_embeddings(batch.prompt_embeds[0], batch.prompt_attention_mask[0]) else: raise ValueError("prompt_attention_mask is None") - # Prepare batch data for Parquet dataset batch_data: list[dict[str, Any]] = [] for key in dataclasses.fields(batch): @@ -175,7 +173,6 @@ class ParquetDatasetSaver: value = value.cpu().numpy() setattr(batch, key.name, value) - # Create record for Parquet dataset records = self.create_records_from_batch(batch) batch_data.extend(records) @@ -199,12 +196,10 @@ class ParquetDatasetSaver: assert isinstance(prompt_attention_mask, torch.Tensor) assert prompt_embeds.shape[0] == prompt_attention_mask.shape[0] - # Get sequence lengths from attention masks (number of 1s) seq_lens = prompt_attention_mask.sum(dim=1) non_padded_embeds = [] - # Process each item in the batch for i in range(prompt_embeds.size(0)): seq_len = seq_lens[i].item() # Slice the embeddings and masks to keep only non-padding parts @@ -224,7 +219,6 @@ class ParquetDatasetSaver: if self._writer is None: return _ = self._writer.flush(write_remainder=write_remainder) - # Reset processed sample count modulo samples_per_file remainder = self.num_processed_samples % self.samples_per_file self.num_processed_samples = 0 if write_remainder else remainder @@ -259,7 +253,6 @@ def build_dataset(preprocess_config: PreprocessConfig, split: str, validator: Ca dataset = dataset.filter(validator) dataset = dataset.shard(num_shards=get_world_size(), index=get_world_rank()) - # add video column def add_video_column(item: dict[str, Any]) -> dict[str, Any]: item["video"] = os.path.join(video_folder, item["name"]) return item diff --git a/fastvideo/workflow/workflow_base.py b/fastvideo/workflow/workflow_base.py index c89fd6209..f8189cd4f 100644 --- a/fastvideo/workflow/workflow_base.py +++ b/fastvideo/workflow/workflow_base.py @@ -47,9 +47,6 @@ class WorkflowBase(ABC): """ self.fastvideo_args = fastvideo_args - # TODO: pipeline_config should be: dict[str, PipelineConfig] - # pipeline_type should be included in the PipelineConfig - # pipeline_config[pipeline_name] = (pipeline_type, fastvideo_args) self._pipeline_configs: dict[str, tuple[PipelineType, FastVideoArgs]] = {} self._pipelines: dict[str, ComposedPipelineBase] = {} self._components: dict[str, Any] = {} diff --git a/scripts/benchmarks/minimax_h3_pro6000/app.py b/scripts/benchmarks/minimax_h3_pro6000/app.py index 5b734d986..f963d8237 100644 --- a/scripts/benchmarks/minimax_h3_pro6000/app.py +++ b/scripts/benchmarks/minimax_h3_pro6000/app.py @@ -540,7 +540,6 @@ def main(step: str = "all", ladder: str = "base"): print("RESULT", json.dumps(result)) -# --- Headline benchmark (480p 5 s, 768p 10 s) on 1 / 4 / 8 GPUs from a FastVideo HF repo. --- # MODAL_PROFILE=aryan5v modal run --detach app.py::headline --repo FastVideo/ --profile h3_dit_ffn --gpus 1,4,8 HERE = pathlib.Path(__file__).resolve().parent headline_image = (image.add_local_file(HERE / "bench_headline.py", "/root/bench_headline.py") diff --git a/scripts/checkpoint_conversion/convert_gen3c_to_fastvideo.py b/scripts/checkpoint_conversion/convert_gen3c_to_fastvideo.py index 2cc24ca34..c2b786e77 100644 --- a/scripts/checkpoint_conversion/convert_gen3c_to_fastvideo.py +++ b/scripts/checkpoint_conversion/convert_gen3c_to_fastvideo.py @@ -121,7 +121,6 @@ def apply_mapping(key: str) -> str | None: if key.startswith(pattern): return None - # Apply mapping patterns for pattern, replacement in PARAM_NAMES_MAPPING.items(): if re.match(pattern, key): return re.sub(pattern, replacement, key) @@ -156,7 +155,6 @@ def load_checkpoint(path: Path, mmap: bool = True) -> dict[str, torch.Tensor]: checkpoint = torch.load(path, **load_kwargs) - # Handle different checkpoint formats if isinstance(checkpoint, dict): for key in ("state_dict", "model_state_dict", "model", "ema"): if key in checkpoint: @@ -254,7 +252,6 @@ def convert_weights( unmapped = [] skipped = [] - # Get all keys first (so we can delete as we go) keys = list(state_dict.keys()) total = len(keys) @@ -310,13 +307,11 @@ def write_component( component_dir = output_dir / name component_dir.mkdir(parents=True, exist_ok=True) - # Save weights output_file = component_dir / "model.safetensors" save_file(weights, str(output_file)) print(f"Saved {name} weights to {output_file}") print(f" {len(weights)} tensors, {sum(t.numel() for t in weights.values()):,} parameters") - # Save config if config is not None: config_path = component_dir / "config.json" with config_path.open("w", encoding="utf-8") as f: @@ -570,7 +565,6 @@ Examples: args = parser.parse_args() - # Validate arguments if args.components_only: if args.output is None: raise ValueError("--output is required with --components-only") @@ -602,7 +596,6 @@ Examples: print(f"Done. Components added to {output_dir}") return - # Parse dtype dtype_map = { "fp32": torch.float32, "fp16": torch.float16, @@ -610,7 +603,6 @@ Examples: } target_dtype = dtype_map[args.dtype] - # Get checkpoint path if args.download: checkpoint_path = download_checkpoint( repo_id=args.download, @@ -622,7 +614,6 @@ Examples: if not checkpoint_path.exists(): raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}") - # Load checkpoint with memory-mapped loading print(f"Loading checkpoint from {checkpoint_path}...") use_mmap = not args.no_mmap state_dict = load_checkpoint(checkpoint_path, mmap=use_mmap) @@ -637,7 +628,6 @@ Examples: total_params = sum(p.numel() for p in state_dict.values()) print(f"Total parameters: {total_params:,} ({total_params / 1e9:.2f}B)") - # Convert weights (memory-efficient mode) print(f"\nConverting weights to {args.dtype}...") converted, unmapped, skipped = convert_weights( state_dict, @@ -662,14 +652,12 @@ Examples: if len(unmapped) > 20: print(f" ... and {len(unmapped) - 20} more") - # Write output output_dir = Path(args.output) if output_dir.exists() and not args.force: raise FileExistsError(f"Output directory exists: {output_dir}. Use --force to overwrite.") output_dir.mkdir(parents=True, exist_ok=True) - # Write transformer weights print(f"\nSaving converted weights...") transformer_config = build_transformer_config() write_component(output_dir, "transformer", converted, transformer_config) @@ -678,7 +666,6 @@ Examples: del converted gc.collect() - # Write model_index.json model_index = build_model_index() model_index_path = output_dir / "model_index.json" with model_index_path.open("w", encoding="utf-8") as f: diff --git a/scripts/checkpoint_conversion/create_hf_repo.py b/scripts/checkpoint_conversion/create_hf_repo.py index d7ecdc78d..25e4eda85 100644 --- a/scripts/checkpoint_conversion/create_hf_repo.py +++ b/scripts/checkpoint_conversion/create_hf_repo.py @@ -87,7 +87,6 @@ def main(): print(f"Removing existing weight file: {ew}") os.remove(ew) - # Copy new weights for src in source_weights: filename = os.path.basename(src) dst = os.path.join(target_component_dir, filename) @@ -108,7 +107,6 @@ def main(): print(f"Uploading to Hugging Face Hub: {target_repo_id}...") try: - # Ensure the repo exists create_repo(target_repo_id, exist_ok=True) upload_folder(repo_id=target_repo_id, folder_path=args.local_dir, commit_message=args.hub_commit_message) diff --git a/scripts/checkpoint_conversion/wan_to_diffusers.py b/scripts/checkpoint_conversion/wan_to_diffusers.py index 2f6e1a626..08866cc73 100644 --- a/scripts/checkpoint_conversion/wan_to_diffusers.py +++ b/scripts/checkpoint_conversion/wan_to_diffusers.py @@ -1,5 +1,4 @@ from huggingface_hub import save_torch_state_dict, load_state_dict_from_file -# from safetensors import safetensors from safetensors.torch import save_file import torch import re @@ -84,7 +83,6 @@ for k, v in state_dict.items(): new_key = re.sub(pattern, replacement, k) break # Stop at the first match else: - # print(f"No match found for {k}") raise ValueError(f"No match found for {k}") new_state_dict[new_key] = v if "norm_added_k" in new_key: diff --git a/scripts/dataset_preparation/generate_validation_json.py b/scripts/dataset_preparation/generate_validation_json.py index e640860ae..1273935c7 100644 --- a/scripts/dataset_preparation/generate_validation_json.py +++ b/scripts/dataset_preparation/generate_validation_json.py @@ -5,7 +5,6 @@ import random def generate_merged_validation_json(input_dir, output_file): - # read in video2caption.json with open(os.path.join(input_dir, "video2caption_replace.json"), "r") as f: video2caption = json.load(f) @@ -30,10 +29,8 @@ def generate_merged_validation_json(input_dir, output_file): } validation_data.append(validation_entry) - # Create the final validation structure validation_json = {"data": validation_data} - # Write the validation JSON to the output file with open(output_file, "w") as f: json.dump(validation_json, f, indent=2) diff --git a/scripts/dataset_preparation/prepare_json_file.py b/scripts/dataset_preparation/prepare_json_file.py index 6b1435890..00b0cdb11 100644 --- a/scripts/dataset_preparation/prepare_json_file.py +++ b/scripts/dataset_preparation/prepare_json_file.py @@ -14,7 +14,6 @@ from fastvideo.eval.io.video import probe_video_fps def get_video_info(video_path): """Get video information.""" - # Read video tensor (T, C, H, W) video_tensor = read_video_tchw(str(video_path)) num_frames = video_tensor.shape[0] @@ -23,7 +22,6 @@ def get_video_info(video_path): fps = probe_video_fps(video_path) or 0 duration = num_frames / fps if fps > 0 else 0 - # Extract name _, _, videos_dir, video_name = str(video_path).split("/") return { @@ -43,7 +41,6 @@ def prepare_dataset_json(folder_path, output_name="videos2caption.json", num_wor """Prepare dataset information from a folder containing videos and prompt.txt.""" folder_path = Path(folder_path) - # Read prompt file prompt_file = folder_path / "prompt.txt" if not prompt_file.exists(): raise FileNotFoundError(f"prompt.txt not found in {folder_path}") @@ -51,7 +48,6 @@ def prepare_dataset_json(folder_path, output_name="videos2caption.json", num_wor with open(prompt_file) as f: prompts = [line.strip() for line in f.readlines() if line.strip()] - # Read videos file videos_file = folder_path / "videos.txt" if not videos_file.exists(): raise FileNotFoundError(f"videos.txt not found in {folder_path}") @@ -62,14 +58,12 @@ def prepare_dataset_json(folder_path, output_name="videos2caption.json", num_wor if len(prompts) != len(video_paths): raise ValueError(f"Number of prompts ({len(prompts)}) does not match number of videos ({len(video_paths)})") - # Prepare arguments for multiprocessing process_args = [folder_path / video_path for video_path in video_paths] # Determine number of workers if num_workers is None: num_workers = max(1, cpu_count() - 1) # Leave one CPU free - # Process videos in parallel start_time = time.time() with Pool(num_workers) as pool: results = list( @@ -84,7 +78,6 @@ def prepare_dataset_json(folder_path, output_name="videos2caption.json", num_wor result["cap"] = [prompt] dataset_info.append(result) - # Calculate total processing time total_time = time.time() - start_time total_videos = len(dataset_info) avg_time_per_video = total_time / total_videos if total_videos > 0 else 0 @@ -94,12 +87,10 @@ def prepare_dataset_json(folder_path, output_name="videos2caption.json", num_wor print(f"Total time: {total_time:.2f} seconds") print(f"Average time per video: {avg_time_per_video:.2f} seconds") - # Save to JSON file output_file = folder_path / output_name with open(output_file, 'w') as f: json.dump(dataset_info, f, indent=2) - # Create merge.txt merge_file = folder_path / "merge.txt" with open(merge_file, 'w') as f: f.write(f"{folder_path}/videos,{output_file}\n") diff --git a/scripts/dataset_preparation/resize_videos.py b/scripts/dataset_preparation/resize_videos.py index dd6a6b4ef..a3afeaeff 100644 --- a/scripts/dataset_preparation/resize_videos.py +++ b/scripts/dataset_preparation/resize_videos.py @@ -74,22 +74,18 @@ def process_folder(args): print(f"Found {len(video_files)} videos") print(f"Target: {args.width}x{args.height} at {args.fps}fps") - # Prepare arguments for parallel processing process_args = [(video_file, output_path, args.width, args.height, args.fps) for video_file in video_files] successful = 0 skipped = 0 failed = [] - # Use ProcessPoolExecutor instead of ThreadPoolExecutor with tqdm(total=len(video_files), desc="Converting videos", dynamic_ncols=True) as pbar: - # Use max_workers as specified or default to CPU count max_workers = args.max_workers with ProcessPoolExecutor(max_workers=max_workers) as executor: # Submit all tasks future_to_file = {executor.submit(resize_video, arg): arg[0] for arg in process_args} - # Process completed tasks for future in as_completed(future_to_file): filename, status, message = future.result() if status == "success": diff --git a/scripts/huggingface/download_hf.py b/scripts/huggingface/download_hf.py index f3c37ac5d..cdf5ebfaa 100644 --- a/scripts/huggingface/download_hf.py +++ b/scripts/huggingface/download_hf.py @@ -2,7 +2,6 @@ import argparse from huggingface_hub import hf_hub_download, snapshot_download -# set args for repo_id, local_dir, repo_type, if __name__ == "__main__": parser = argparse.ArgumentParser(description="Download a dataset or model from the Hugging Face Hub") diff --git a/scripts/lora_extraction/lora_inference_comparison.py b/scripts/lora_extraction/lora_inference_comparison.py index 9fc757d75..c11a06dc2 100644 --- a/scripts/lora_extraction/lora_inference_comparison.py +++ b/scripts/lora_extraction/lora_inference_comparison.py @@ -72,7 +72,6 @@ def _validate_adapter(adapter: Optional[str]) -> Optional[str]: if p.stat().st_size == 0: raise ValueError(f"Adapter file is empty: {p}") elif p.is_dir(): - # Check if directory contains at least one .safetensors file safetensors_files = list(p.glob("*.safetensors")) if not safetensors_files: raise ValueError(f"Adapter directory contains no .safetensors files: {p}") @@ -172,11 +171,9 @@ def compute_metrics(output_dir: str, ft_video: str, other_video: str, num_infere loss_fn = lpips.LPIPS(net='alex') - # Load videos vid1, _, _ = tv_io.read_video(ft_video, pts_unit='sec') vid2, _, _ = tv_io.read_video(other_video, pts_unit='sec') - # Normalize to [-1, 1] vid1 = (vid1.float() / 127.5 - 1.0).permute(0, 3, 1, 2) # (T, C, H, W) vid2 = (vid2.float() / 127.5 - 1.0).permute(0, 3, 1, 2) @@ -188,7 +185,6 @@ def compute_metrics(output_dir: str, ft_video: str, other_video: str, num_infere results["mean_lpips"] = sum(lpips_scores) / len(lpips_scores) - # Write LPIPS results import json lpips_file = Path(output_dir) / f"steps{num_inference_steps}_{prompt.replace(' ', '_')[:30]}_lpips.json" with open(lpips_file, 'w') as f: diff --git a/scripts/verify_anyflow_fastvideo_parity.py b/scripts/verify_anyflow_fastvideo_parity.py index 85a734bd8..686b529f9 100644 --- a/scripts/verify_anyflow_fastvideo_parity.py +++ b/scripts/verify_anyflow_fastvideo_parity.py @@ -54,10 +54,8 @@ def banner(msg: str) -> None: print("=" * 80) -# --------------------------------------------------------------------------- # Distributed bootstrap (single-rank). Required by WanTransformer3DModel # which calls get_sp_world_size() and uses ReplicatedLinear. -# --------------------------------------------------------------------------- def init_single_rank() -> None: @@ -79,10 +77,8 @@ def init_single_rank() -> None: print(" single-rank distributed environment + TP/SP/DP groups initialized") -# --------------------------------------------------------------------------- # Translate AnyFlow HF safetensor keys onto FastVideo's WanTransformer3DModel # internal layout, applying the regex from WanVideoArchConfig.param_names_mapping. -# --------------------------------------------------------------------------- def translate_keys( @@ -99,9 +95,6 @@ def translate_keys( return out -# --------------------------------------------------------------------------- -# Build FastVideo WanTransformer3DModel with AnyFlow weights loaded. -# --------------------------------------------------------------------------- def build_fastvideo_transformer(): @@ -133,7 +126,6 @@ def build_fastvideo_transformer(): print(f" Wan transformer built in {time.time() - t0:.1f}s; " f"params: {sum(p.numel() for p in model.parameters())/1e9:.2f}B") - # Load AnyFlow checkpoint. af_path = ANYFLOW_LOCAL / "transformer" / "diffusion_pytorch_model.safetensors" af_raw = load_file(str(af_path), device="cpu") print(f" AnyFlow checkpoint: {len(af_raw)} tensors") @@ -146,9 +138,6 @@ def build_fastvideo_transformer(): return model, len(miss), len(unex) -# --------------------------------------------------------------------------- -# Build AnyFlow reference net. -# --------------------------------------------------------------------------- def build_anyflow_reference(): @@ -166,9 +155,6 @@ def build_anyflow_reference(): return af_net -# --------------------------------------------------------------------------- -# Forward parity test. -# --------------------------------------------------------------------------- def forward_compare(fv_model, af_net): @@ -232,9 +218,6 @@ def forward_compare(fv_model, af_net): return ok -# --------------------------------------------------------------------------- -# Any-step sampling smoke via FlowMapEulerDiscreteScheduler. -# --------------------------------------------------------------------------- def sample_anystep(fv_model):