diff --git a/fastvideo/models/dits/cosmos3.py b/fastvideo/models/dits/cosmos3.py index 074f646de..2a0094dd2 100644 --- a/fastvideo/models/dits/cosmos3.py +++ b/fastvideo/models/dits/cosmos3.py @@ -46,6 +46,7 @@ import torch.nn.functional as F from fastvideo.configs.models.dits.cosmos3 import Cosmos3VideoConfig from fastvideo.layers.layernorm import RMSNorm from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.visual_embedding import timestep_embedding from fastvideo.models.dits.base import BaseDiT EntryClass = ["Cosmos3VFMTransformer"] @@ -320,19 +321,8 @@ class Cosmos3TimestepEmbedder(nn.Module): self.act = nn.SiLU() self.linear_2 = ReplicatedLinear(hidden_size, hidden_size, bias=True) - @staticmethod - def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: - half = dim // 2 - freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( - device=t.device) - args = t[:, None].float() * freqs[None] - embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - if dim % 2: - embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) - return embedding - def forward(self, t: torch.Tensor) -> torch.Tensor: - t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + t_freq = timestep_embedding(t, self.frequency_embedding_size) # Sinusoidal runs in fp32 (timestep_scale makes inputs tiny); cast to the # MLP weight dtype (bf16 at inference; no-op in the fp32 parity tests). t_freq = t_freq.to(self.linear_1.weight.dtype) diff --git a/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py b/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py index 9d989734c..dee2837ae 100644 --- a/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py +++ b/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py @@ -60,15 +60,6 @@ logger = init_logger(__name__) _SYSTEM_PROMPT_IMAGE = "You are a helpful assistant who will generate images from a give prompt." _SYSTEM_PROMPT_VIDEO = "You are a helpful assistant who will generate videos from a give prompt." -# Default video negative prompt (framework / Cosmos quality prompt). -COSMOS3_VIDEO_NEGATIVE_PROMPT = ( - "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, " - "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, " - "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, " - "jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, " - "fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. " - "Overall, the video is of poor quality.") - # =========================================================================== # Special-token resolution (Qwen2 chat tokenizer) @@ -175,15 +166,12 @@ class Cosmos3VisionSpec: """Geometry + conditioning for one vision item in a denoise run. Args: - clean_latent: VAE-encoded conditioning latent ``[C, T, H, W]`` used to - keep condition frames clean (I2V / T2I). May be ``None`` for T2V. condition_frame_indexes: Latent-frame indices kept clean. shape: ``(C, T, H, W)`` of the latent for this item. """ shape: tuple[int, int, int, int] condition_frame_indexes: list[int] - clean_latent: torch.Tensor | None = None @property def numel(self) -> int: @@ -244,6 +232,7 @@ def cosmos3_get_cfg_velocity( timestep_value = float(timestep.reshape(()).item()) noise_x_vision = _split_flat_latent(flat_latent, specs) + device = next(transformer.parameters()).device def _run(token_ids: list[int]) -> torch.Tensor: samples = [ @@ -268,7 +257,6 @@ def cosmos3_get_cfg_velocity( base_fps=base_fps, temporal_compression_factor=temporal_compression_factor, ) - device = next(transformer.parameters()).device out = transformer(**packed.to_dit_kwargs(device=device)) preds = out.get("preds_vision") if preds is None: diff --git a/fastvideo/pipelines/basic/cosmos3/presets.py b/fastvideo/pipelines/basic/cosmos3/presets.py index d4053ebfc..197d27907 100644 --- a/fastvideo/pipelines/basic/cosmos3/presets.py +++ b/fastvideo/pipelines/basic/cosmos3/presets.py @@ -20,7 +20,7 @@ _DENOISE_STAGE = PresetStageSpec( ) # Framework video negative prompt (Cosmos quality prompt). -_COSMOS3_NEGATIVE_PROMPT = ( +COSMOS3_VIDEO_NEGATIVE_PROMPT = ( "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, " "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, " "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, " @@ -42,7 +42,7 @@ COSMOS3_NANO = InferencePreset( "fps": 24, "guidance_scale": 6.0, "num_inference_steps": 35, - "negative_prompt": _COSMOS3_NEGATIVE_PROMPT, + "negative_prompt": COSMOS3_VIDEO_NEGATIVE_PROMPT, }, ) @@ -60,7 +60,7 @@ COSMOS3_NANO_I2V = InferencePreset( "fps": 24, "guidance_scale": 6.0, "num_inference_steps": 35, - "negative_prompt": _COSMOS3_NEGATIVE_PROMPT, + "negative_prompt": COSMOS3_VIDEO_NEGATIVE_PROMPT, }, ) diff --git a/fastvideo/pipelines/stages/cosmos3_stages.py b/fastvideo/pipelines/stages/cosmos3_stages.py index e4b925fdb..b0bf596b2 100644 --- a/fastvideo/pipelines/stages/cosmos3_stages.py +++ b/fastvideo/pipelines/stages/cosmos3_stages.py @@ -30,7 +30,6 @@ from tqdm import tqdm from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( - COSMOS3_VIDEO_NEGATIVE_PROMPT, Cosmos3DenoiseEngine, Cosmos3VisionSpec, _VaeNorm, @@ -39,6 +38,8 @@ from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( cosmos3_vae_decode, cosmos3_vae_encode, ) +from fastvideo.pipelines.basic.cosmos3.presets import ( + COSMOS3_VIDEO_NEGATIVE_PROMPT, ) from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.base import PipelineStage @@ -141,7 +142,6 @@ class Cosmos3DenoisingStage(PipelineStage): spec = Cosmos3VisionSpec( shape=latent_shape, condition_frame_indexes=condition_frame_indexes, - clean_latent=clean_latent, ) # ---- Scheduler timesteps ----