diff --git a/fastgen/configs/experiments/CosmosPredict2/config_dmd2.py b/fastgen/configs/experiments/CosmosPredict2/config_dmd2.py index 408e391..2ea477a 100644 --- a/fastgen/configs/experiments/CosmosPredict2/config_dmd2.py +++ b/fastgen/configs/experiments/CosmosPredict2/config_dmd2.py @@ -28,8 +28,8 @@ def create_config(): # Small resolution for testing: 320x176 video -> 40x22 latent, 93 frames -> 24 latent # Full 720p: [16, 24, 88, 160] (1280x704 @ 93 frames) # config.model.input_shape = [16, 24, 22, 40] # cthw - 256p (320x176 video) - config.model.input_shape = [16, 24, 60, 104] # cthw - 480p, 93 frames - # config.model.input_shape = [16, 24, 88, 160] # cthw - full 720p, 93 frames + # config.model.input_shape = [16, 24, 60, 104] # cthw - 480p, 93 frames + config.model.input_shape = [16, 24, 88, 160] # cthw - full 720p, 93 frames # Network and discriminator config.model.net = CosmosPredict2_2B_Config @@ -44,7 +44,8 @@ def create_config(): config.model.gan_use_same_t_noise = True config.model.fake_score_pred_type = "x0" config.model.student_sample_type = "ode" - config.model.guidance_scale = 3.0 + # Official 2B DMD2 teacher guidance scale + config.model.guidance_scale = 4.0 config.model.pretrained_model_path = f"{CKPT_ROOT_DIR}/cosmos_predict2/Cosmos-Predict2.5-2B/base/post-trained/81edfebe-bd6a-4039-8c1d-737df1a790bf_ema_bf16.pt" # Timestep sampling diff --git a/fastgen/configs/experiments/CosmosPredict2/config_dmd2_14b.py b/fastgen/configs/experiments/CosmosPredict2/config_dmd2_14b.py index bea7329..a1d403d 100644 --- a/fastgen/configs/experiments/CosmosPredict2/config_dmd2_14b.py +++ b/fastgen/configs/experiments/CosmosPredict2/config_dmd2_14b.py @@ -3,6 +3,8 @@ """DMD2 config for Cosmos-Predict2.5-14B model.""" +import copy + import fastgen.configs.experiments.CosmosPredict2.config_dmd2 as config_dmd2_base from fastgen.configs.discriminator import Discriminator_CosmosPredict2_14B_Config from fastgen.configs.net import CosmosPredict2_14B_Config, CosmosPredict2_14B_Aggressive_Config, CKPT_ROOT_DIR @@ -29,6 +31,7 @@ def create_config(): config.model.pretrained_model_path = f"{CKPT_ROOT_DIR}/cosmos_predict2/Cosmos-Predict2.5-14B/base/post-trained/e21d2a49-4747-44c8-ba44-9f6f9243715f_ema_bf16.pt" # Dataloader img_size = (W, H) = (latent_W * 8, latent_H * 8) + config.dataloader_train = copy.deepcopy(config.dataloader_train) config.dataloader_train.img_size = (config.model.input_shape[-1] * 8, config.model.input_shape[-2] * 8) config.dataloader_train.sequence_length = (config.model.input_shape[1] - 1) * 4 + 1 diff --git a/fastgen/configs/experiments/CosmosPredict2/config_dmd2_v2w.py b/fastgen/configs/experiments/CosmosPredict2/config_dmd2_v2w.py index c428953..d084e8b 100644 --- a/fastgen/configs/experiments/CosmosPredict2/config_dmd2_v2w.py +++ b/fastgen/configs/experiments/CosmosPredict2/config_dmd2_v2w.py @@ -3,6 +3,8 @@ """DMD2 config for Cosmos-Predict2.5-2B video2world model.""" +import copy + import fastgen.configs.experiments.CosmosPredict2.config_dmd2 as config_dmd2_base @@ -10,6 +12,7 @@ def create_config(): config = config_dmd2_base.create_config() # Network for v2w + config.model.net = copy.deepcopy(config.model.net) config.model.net.is_video2world = True config.model.net.num_conditioning_frames = 1 diff --git a/fastgen/configs/experiments/CosmosPredict2/config_sft.py b/fastgen/configs/experiments/CosmosPredict2/config_sft.py index 6377bd6..2a62787 100644 --- a/fastgen/configs/experiments/CosmosPredict2/config_sft.py +++ b/fastgen/configs/experiments/CosmosPredict2/config_sft.py @@ -14,7 +14,7 @@ def create_config(): config.trainer.logging_iter = 100 config.trainer.save_ckpt_iter = 500 - config.model.net_optimizer.lr = 1e-5 + config.model.net_optimizer.lr = 5e-6 config.model.sample_t_cfg.time_dist_type = "uniform" config.model.sample_t_cfg.min_t = 0.001 @@ -27,13 +27,13 @@ def create_config(): # Small resolution for testing: 320x176 video -> 40x22 latent, 21 frames -> 6 latent # Full 720p: [16, 24, 88, 160] (1280x704 @ 93 frames) # config.model.input_shape = [16, 24, 24, 40] # cthw - 256p - # config.model.input_shape = [16, 21, 60, 104] # cthw - 480p, 81 frames - config.model.input_shape = [16, 24, 60, 104] # cthw - 480p - # config.model.input_shape = [16, 24, 88, 160] # cthw - full 720p + # config.model.input_shape = [16, 24, 60, 104] # cthw - 480p + config.model.input_shape = [16, 24, 88, 160] # cthw - full 720p config.model.net = CosmosPredict2_2B_Config config.model.pretrained_model_path = f"{CKPT_ROOT_DIR}/cosmos_predict2/Cosmos-Predict2.5-2B/base/post-trained/81edfebe-bd6a-4039-8c1d-737df1a790bf_ema_bf16.pt" - config.model.guidance_scale = 3.0 + # Official post-trained 2B/14B inference uses guidance scale 7 + config.model.guidance_scale = 5.0 config.model.student_sample_steps = 35 config.dataloader_train = VideoLoaderConfig diff --git a/fastgen/configs/experiments/CosmosPredict2/config_sft_v2w.py b/fastgen/configs/experiments/CosmosPredict2/config_sft_v2w.py index 5c38adc..211b3a3 100644 --- a/fastgen/configs/experiments/CosmosPredict2/config_sft_v2w.py +++ b/fastgen/configs/experiments/CosmosPredict2/config_sft_v2w.py @@ -3,6 +3,8 @@ """Configs for SFT on Cosmos-Predict2.5-2B video2world model.""" +import copy + import fastgen.configs.experiments.CosmosPredict2.config_sft as config_sft_base @@ -10,8 +12,11 @@ def create_config(): config = config_sft_base.create_config() # Network for v2w + config.model.net = copy.deepcopy(config.model.net) config.model.net.is_video2world = True config.model.net.num_conditioning_frames = 1 + # Official post-trained Video2World inference uses guidance scale 8 + config.model.guidance_scale = 5.0 config.log_config.group = "cosmos_predict2_sft_v2w" diff --git a/fastgen/configs/net.py b/fastgen/configs/net.py index 7833c0b..f1aa667 100644 --- a/fastgen/configs/net.py +++ b/fastgen/configs/net.py @@ -277,7 +277,10 @@ num_blocks=28, num_heads=16, sac_config=L(SACConfig)(mode=CheckpointMode.BLOCK_WISE), - fps=24, + # Checkpoint-specific override: the released post-trained 2B model uses + # 16 FPS with temporal RoPE FPS modulation disabled. + fps=16, + rope_enable_fps_modulation=False, is_video2world=False, num_conditioning_frames=1, enable_logvar_linear=False, # Enable logvar_linear for sCM-like models diff --git a/fastgen/networks/cosmos_predict2/network.py b/fastgen/networks/cosmos_predict2/network.py index 0fd14cc..b6215f7 100644 --- a/fastgen/networks/cosmos_predict2/network.py +++ b/fastgen/networks/cosmos_predict2/network.py @@ -767,6 +767,7 @@ def __init__( use_wan_fp32_strategy: bool = True, # FPS for temporal position embeddings fps: float = 24.0, + rope_enable_fps_modulation: bool = True, # Video2world (image-to-video) mode settings is_video2world: bool = False, num_conditioning_frames: int = 1, @@ -807,6 +808,7 @@ def __init__( rope_h_extrapolation_ratio=rope_h_extrapolation_ratio, rope_w_extrapolation_ratio=rope_w_extrapolation_ratio, rope_t_extrapolation_ratio=rope_t_extrapolation_ratio, + rope_enable_fps_modulation=rope_enable_fps_modulation, use_wan_fp32_strategy=use_wan_fp32_strategy, ) @@ -1097,13 +1099,12 @@ def sample( neg_condition: Optional[Any] = None, guidance_scale: Optional[float] = 5.0, num_steps: int = 50, - shift: float = 5.0, # Cosmos default skip_layers: Optional[List[int]] = None, skip_layers_start_percent: float = 0.0, fps: Optional[torch.Tensor] = None, conditioning_latents: Optional[torch.Tensor] = None, num_conditioning_frames: int = 1, - conditional_frame_timestep: float = 0.0, + conditional_frame_timestep: float = 1e-4, denoise_replace_gt_frames: bool = True, **kwargs, ) -> torch.Tensor: @@ -1125,18 +1126,18 @@ def sample( noise: The noisy latents to start from, shape (B, C, T, H, W). condition: The conditioning embeddings. neg_condition: The negative/unconditional embeddings for CFG. - guidance_scale: Classifier-free guidance scale. Values > 1.0 enable CFG. - num_steps: The number of sampling steps. - shift: Noise schedule shift parameter. + guidance_scale: Conventional classifier-free guidance scale. + Values greater than 1.0 enable CFG. + num_steps: Number of intervals in the official Karras schedule. + The sampler evaluates ``num_steps + 1`` timesteps. skip_layers: List of transformer layers to skip (for skip-layer guidance). skip_layers_start_percent: Percentage of steps before starting to skip layers. fps: Frames per second tensor for temporal conditioning. conditioning_latents: Latent frames to condition on for video2world mode, shape (B, C, T, H, W). If provided, enables video2world mode. num_conditioning_frames: Number of frames to use for conditioning (default 1). - conditional_frame_timestep: Timestep value for conditioning frames (default 0.0). - Use 0.0 to indicate clean frames with no noise, or a negative value - to disable the special conditioning-frame timestep. + conditional_frame_timestep: Normalized conditioning-frame timestep + forwarded to ``forward()`` (default 1e-4). denoise_replace_gt_frames: Whether to replace velocity for conditioning frames with analytical velocity (noise - gt_frames). Default True. @@ -1145,17 +1146,19 @@ def sample( """ assert self.schedule_type == "rf", f"{self.schedule_type} is not supported" - # Set timesteps with shift parameter (matching official Cosmos inference) + # Match official Cosmos Predict2.5 inference. Diffusers uses the + # configured sigma bounds for its Karras conversion; the official + # `num_steps` denotes intervals, hence `num_steps + 1` ramp points. if self.sample_scheduler is None: self.sample_scheduler = UniPCMultistepScheduler( num_train_timesteps=1000, prediction_type="flow_prediction", use_flow_sigmas=True, - flow_shift=shift, + use_karras_sigmas=True, + sigma_min=0.01, + sigma_max=200.0, ) - else: - self.sample_scheduler.config.flow_shift = shift - self.sample_scheduler.set_timesteps(num_inference_steps=num_steps, device=noise.device) + self.sample_scheduler.set_timesteps(num_inference_steps=num_steps + 1, device=noise.device) timesteps = self.sample_scheduler.timesteps # Initialize latents with proper scaling based on the initial timestep @@ -1211,7 +1214,7 @@ def sample( # Store initial noise for velocity replacement initial_noise = latents.clone() - for idx, timestep in tqdm(enumerate(timesteps), total=num_steps, desc="Sampling"): + for timestep in tqdm(timesteps, total=len(timesteps), desc="Sampling"): # Normalize timestep to [0, 1] range t = (timestep / self.sample_scheduler.config.num_train_timesteps).expand(latents.shape[0]) t = self.noise_scheduler.safe_clamp(t, min=self.noise_scheduler.min_t, max=self.noise_scheduler.max_t).to( @@ -1264,14 +1267,8 @@ def sample( gt_velocity = initial_noise - conditioning_latents_full velocity_pred = gt_velocity * condition_mask_C + velocity_pred * (1 - condition_mask_C) - # Scheduler step (use model_input for video2world since that's what velocity was computed on) - sample_input = model_input if video2world_mode else latents - latents = self.sample_scheduler.step(velocity_pred, timestep, sample_input, return_dict=False)[0] - - # Preserve conditioning frames after each step to prevent numerical drift - if video2world_mode: - v2w_condition = {"conditioning_latents": conditioning_latents, "condition_mask": condition_mask} - latents = self.preserve_conditioning(latents, v2w_condition) + # Keep clean frames in the DiT input while UniPC evolves raw latents. + latents = self.sample_scheduler.step(velocity_pred, timestep, latents, return_dict=False)[0] return latents @@ -1316,9 +1313,10 @@ def forward( fps: Frames per second tensor padding_mask: Padding mask tensor skip_layers: List of block indices to skip during forward pass - conditional_frame_timestep: Timestep value for conditioning frames (default 0.0). - Use 0.0 to indicate clean frames with no noise, or a negative value - to disable the special conditioning-frame timestep. + conditional_frame_timestep: Normalized timestep for conditioning frames. + ``forward()`` defaults to 0.0; ``sample()`` passes 1e-4, equivalent + to the official raw timestep 0.1 after its 0.001 network scale. + A negative value disables the special conditioning-frame timestep. Returns: Depending on the arguments: diff --git a/scripts/README.md b/scripts/README.md index a8d62a2..4c49159 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -82,12 +82,9 @@ python scripts/inference/video_model_inference.py \ ### Video2World (Cosmos) ```bash python scripts/inference/video_model_inference.py \ - --config fastgen/configs/experiments/CosmosPredict2/config_sft.py \ - --do_student_sampling False --num_steps 35 --fps 24 \ - --neg_prompt_file scripts/inference/prompts/negative_prompt_cosmos.txt \ - --input_image_file scripts/inference/prompts/source_image_paths.txt \ - --num_conditioning_frames 1 \ - - model.guidance_scale=5.0 model.net.is_video2world=True model.input_shape="[16, 24, 88, 160]" + --config fastgen/configs/experiments/CosmosPredict2/config_sft_v2w.py \ + --do_student_sampling False --num_steps 35 \ + --neg_prompt_file scripts/inference/prompts/negative_prompt_cosmos.txt ``` ### Causal / Autoregressive diff --git a/scripts/inference/video_model_inference.py b/scripts/inference/video_model_inference.py index 13cc85b..f18ee4a 100644 --- a/scripts/inference/video_model_inference.py +++ b/scripts/inference/video_model_inference.py @@ -28,16 +28,15 @@ PYTHONPATH=$(pwd) FASTGEN_OUTPUT_ROOT='FASTGEN_OUTPUT' torchrun --nproc_per_node=1 --standalone \\ scripts/inference/video_model_inference.py --do_student_sampling False \\ --source_video_file scripts/inference/prompts/source_video_paths.txt \\ - --config fastgen/configs/experiments/WanV2V/config_sft_latent.py \\ + --config fastgen/configs/experiments/WanV2V/config_sft.py \\ - trainer.seed=1 trainer.ddp=True model.guidance_scale=5.0 log_config.name=vace_wan_inference # Video2World: Cosmos Predict2 PYTHONPATH=$(pwd) FASTGEN_OUTPUT_ROOT='FASTGEN_OUTPUT' torchrun --nproc_per_node=1 --standalone \\ - scripts/inference/video_model_inference.py --do_student_sampling False \\ - --input_image_file scripts/inference/prompts/source_image_paths.txt --num_conditioning_frames 1 \\ - --config fastgen/configs/experiments/CosmosPredict2/config_sft.py \\ - - trainer.seed=1 trainer.ddp=True model.guidance_scale=5.0 model.net.is_video2world=True \\ - log_config.name=cosmos_v2w_inference + scripts/inference/video_model_inference.py --do_student_sampling False --num_steps 35 \\ + --neg_prompt_file scripts/inference/prompts/negative_prompt_cosmos.txt \\ + --config fastgen/configs/experiments/CosmosPredict2/config_sft_v2w.py \\ + - trainer.ddp=True log_config.name=cosmos_v2w_inference # Eval with skip-layer guidance (SLG) PYTHONPATH=$(pwd) FASTGEN_OUTPUT_ROOT='FASTGEN_OUTPUT' torchrun --nproc_per_node=1 --standalone \\ @@ -721,7 +720,7 @@ def main(args, config: BaseConfig): "--fps", default=16, type=int, - help="Frames per second for saved video and model temporal encoding (default: 16, matches Wan base_fps)", + help="Frames per second for saved video and model temporal encoding (default: 16)", ) parser.add_argument( "--save_high_quality", @@ -801,7 +800,7 @@ def main(args, config: BaseConfig): "--num_conditioning_frames", type=int, default=1, - help="Number of latent frames to condition on for I2V mode (default: 1).", + help="Number of latent frames to condition on for I2V/Video2World mode (default: 1).", ) args = parse_args(parser)