nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[bugfix]: kandinsky5 stages: official rng order, attn metadata, tensor resize, decode dedup

Browse Source 
Match the official kandinskylab/kandinsky-5 RNG order for I2V seed parity:
image encoding now runs after latent preparation and places the conditioning
latent itself, so the initial noise is the seeded generator's first draw
(upstream leaves the image-latent sample unseeded; drawing it second from the
same generator keeps FastVideo deterministic). The denoising stage now sets
the forward context with NABLA attention metadata per step — which also
routes the dense path through LocalAttention instead of the no-context SDPA
fallback. Resize tensor conditioning images to the requested resolution like
the PIL path. Collapse Kandinsky5DecodingStage to a channels-first permute +
super().forward(), regaining the base class's MPS handling.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Will Lin
2026-07-07 14:35:48 -07:00
committed by SolitaryThinker
co-authored by Claude Fable 5
parent 4bc87f1934
commit 1d2f329a54
2 changed files with 49 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/pipelines/basic/kandinsky5/kandinsky5_i2v_pipeline.py
+8 -5
View File
@@ -46,11 +46,6 @@ class Kandinsky5I2VPipeline(ComposedPipelineBase):
),
)
self.add_stage(
stage_name="image_encoding_stage",
stage=Kandinsky5ImageEncodingStage(vae=self.get_module("vae")),
)
self.add_stage(
stage_name="timestep_preparation_stage",
stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
@@ -64,6 +59,14 @@ class Kandinsky5I2VPipeline(ComposedPipelineBase):
),
)
# Runs AFTER latent preparation so the initial noise is the seeded
# generator's first draw (official kandinsky-5 RNG order); this stage
# then samples the image latent and places it into the latents.
self.add_stage(
stage_name="image_encoding_stage",
stage=Kandinsky5ImageEncodingStage(vae=self.get_module("vae")),
)
self.add_stage(
stage_name="denoising_stage",
stage=Kandinsky5DenoisingStage(
fastvideo/pipelines/stages/kandinsky5.py
+41 -38
View File
@@ -10,8 +10,10 @@ import torch
from diffusers.utils.torch_utils import randn_tensor
from tqdm.auto import tqdm
from fastvideo.attention.backends.nabla import NablaAttentionMetadataBuilder
from fastvideo.distributed import get_local_torch_device
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.forward_context import set_forward_context
from fastvideo.logger import init_logger
from fastvideo.models.loader.component_loader import TransformerLoader, VAELoader
from fastvideo.models.vaes.common import ParallelTiledVAE
@@ -108,15 +110,9 @@ class Kandinsky5LatentPreparationStage(PipelineStage):
)
latents = torch.cat([latents, cond, cond_mask], dim=-1)
# I2V: place the encoded image in the first frame and, when the
# transformer expects visual conditioning, in the cond/mask channels.
if batch.image_latent is not None:
image_latent = batch.image_latent.to(device=device, dtype=dtype)
latents[:, 0:1, :, :, :num_channels] = image_latent
if visual_cond:
latents[:, 0:1, :, :, num_channels:2 * num_channels] = image_latent
latents[:, 0:1, :, :, 2 * num_channels:] = 1.0
# I2V image conditioning is placed by Kandinsky5ImageEncodingStage,
# which runs AFTER this stage so the initial noise is the generator's
# first draw (matching the official kandinskylab/kandinsky-5 order).
batch.latents = latents
batch.raw_latent_shape = (
batch_size,
@@ -316,9 +312,18 @@ class Kandinsky5DenoisingStage(PipelineStage):
continue
t_expand = timestep.unsqueeze(0).repeat(latents.shape[0]).to(device=device, dtype=target_dtype)
attn_metadata = None
if sparse_params is not None:
attn_metadata = NablaAttentionMetadataBuilder().build(
current_timestep=i,
sta_mask=sparse_params["sta_mask"],
P=sparse_params["P"],
visual_shape=sparse_params["visual_shape"],
)
autocast_ctx = (torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled)
if device.type == "cuda" else contextlib.nullcontext())
with autocast_ctx:
with set_forward_context(current_timestep=i, attn_metadata=attn_metadata,
forward_batch=batch), autocast_ctx:
pred_velocity = self.transformer(
hidden_states=latents.to(dtype=target_dtype),
encoder_hidden_states=prompt_embeds,
@@ -381,36 +386,13 @@ class Kandinsky5DecodingStage(DecodingStage):
super().__init__(vae=vae, pipeline=pipeline)
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
pipeline = self.pipeline() if self.pipeline else None
if not fastvideo_args.model_loaded["vae"]:
loader = VAELoader()
self.vae = loader.load(fastvideo_args.model_paths["vae"], fastvideo_args)
if pipeline:
pipeline.add_module("vae", self.vae)
fastvideo_args.model_loaded["vae"] = True
if batch.latents is None:
raise ValueError("latents must be available before Kandinsky5 decoding.")
if fastvideo_args.output_type == "latent":
frames = batch.latents.permute(0, 4, 1, 2, 3).contiguous()
else:
frames = self.decode(batch.latents.permute(0, 4, 1, 2, 3).contiguous(), fastvideo_args)
# Trajectory latents are recorded channels-first by the denoising stage.
if batch.return_trajectory_decoded:
batch.trajectory_decoded = []
assert batch.trajectory_latents is not None, "batch should have trajectory latents"
for idx in range(batch.trajectory_latents.shape[1]):
cur_latent = batch.trajectory_latents[:, idx]
logger.info("decoding trajectory latent for timestep: %s", batch.trajectory_timesteps[idx])
decoded_frames = self.decode(cur_latent, fastvideo_args)
batch.trajectory_decoded.append(decoded_frames.cpu().float())
batch.output = frames.to(torch.float32)
if fastvideo_args.vae_cpu_offload:
self.vae.to("cpu")
return batch
# Kandinsky5 latents are channels-last [B, T, H, W, C]; the base stage
# (and the trajectory latents recorded by the denoising stage) work
# channels-first.
batch.latents = batch.latents.permute(0, 4, 1, 2, 3).contiguous()
return super().forward(batch, fastvideo_args)
class Kandinsky5ImageEncodingStage(EncodingStage):
@@ -433,6 +415,13 @@ class Kandinsky5ImageEncodingStage(EncodingStage):
"assuming range [0, 1] and normalizing to [-1, 1]. "
"Pass a [-1, 1] tensor with negative values to skip normalization.")
image = normalize(image) # [0, 1] -> [-1, 1]
if image.ndim == 3:
image = image.unsqueeze(0)
if image.shape[-2:] != (height, width):
image = torch.nn.functional.interpolate(image.float(),
size=(height, width),
mode="bilinear",
antialias=True)
return image
@torch.no_grad()
@@ -480,6 +469,18 @@ class Kandinsky5ImageEncodingStage(EncodingStage):
# [B, C, 1, H, W] -> [B, 1, H, W, C] to match channels-last latents
batch.image_latent = image_latent.permute(0, 2, 3, 4, 1).contiguous()
# Place the conditioning latent into the prepared latents: frame 0 of
# the main channels and, when the transformer expects visual
# conditioning, the cond/mask channels.
latents = batch.latents
image_latent = batch.image_latent.to(device=latents.device, dtype=latents.dtype)
num_channels = image_latent.shape[-1]
latents[:, 0:1, :, :, :num_channels] = image_latent
if latents.shape[-1] > num_channels:
latents[:, 0:1, :, :, num_channels:2 * num_channels] = image_latent
latents[:, 0:1, :, :, 2 * num_channels:] = 1.0
batch.latents = latents
if fastvideo_args.vae_cpu_offload:
vae.to("cpu")
return batch
@@ -489,6 +490,8 @@ class Kandinsky5ImageEncodingStage(EncodingStage):
result.add_check("pil_image", batch.pil_image, V.not_none)
result.add_check("height", batch.height, V.positive_int)
result.add_check("width", batch.width, V.positive_int)
# This stage runs after latent preparation and writes into its output.
result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
return result
def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: