LTX2 Latent Optimization issue

[asafjo23]

Hi,

I'm trying to capture the model prompt cross attention in order to apply some latent optimization techniques during inference, but for some reason I'm getting latents.grad as None no matter what I'm trying. I'll add small code snippets to describe what I'm trying to do:

class Optimizer: def __init__( self, loss_fn: LayoutLoss, num_refinements: int = 3, lr_start: float = 0.01, lr_end: float = 0.05, betas: tuple[float, float] = (0.4, 0.9), weight_decay: float = 0.0, ): self.loss_fn = loss_fn self.num_refinements = num_refinements self.lr_start = lr_start self.lr_end = lr_end self.betas = betas self.weight_decay = weight_decay def optimize( self, transformer: LTX2VideoTransformer3DModel, latents: torch.Tensor, audio_latents: torch.Tensor, prompt_embeds: torch.Tensor, audio_prompt_embeds: torch.Tensor, timestep: torch.Tensor, attention_mask: torch.Tensor, num_frames: int, height: int, width: int, fps: float, audio_num_frames: int, video_coords: torch.Tensor, audio_coords: torch.Tensor, attention_kwargs: Dict[str, Any], store: AttentionStore, progress_bar: tqdm.tqdm, ) -> torch.Tensor: latents = latents.clone().detach() latents = latents.to(transformer.dtype) optimizer = torch.optim.AdamW( [latents], lr=self.lr_start, betas=self.betas, weight_decay=self.weight_decay, ) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=self.num_refinements, eta_min=self.lr_end ) transformer.zero_grad(set_to_none=True) first_loss = None with torch.enable_grad(): for i in range(self.num_refinements): latents = latents.requires_grad_(True) store.reset() latent_model_input = latents.to(transformer.dtype) _ = transformer( hidden_states=latent_model_input, audio_hidden_states=audio_latents, encoder_hidden_states=prompt_embeds, audio_encoder_hidden_states=audio_prompt_embeds, timestep=timestep, encoder_attention_mask=attention_mask, audio_encoder_attention_mask=attention_mask, num_frames=num_frames, height=height, width=width, fps=fps, audio_num_frames=audio_num_frames, video_coords=video_coords, audio_coords=audio_coords, attention_kwargs=attention_kwargs, return_dict=False, )[0] attn = store.get_avg_attention().unsqueeze(0) loss = self.loss_fn(attn) # Backward loss.backward() # FIX 4: Verify gradients exist before stepping if latents.grad is None: print(f"WARNING: latents.grad is None at iteration {i+1}!") print(" Gradient flow is broken. Check:") print(" 1. AttentionStore doesn't use .clone()") print(" 2. No dtype conversion breaks the computation graph") print(" 3. Gradient checkpointing is disabled") break # Only step if we have gradients optimizer.step() scheduler.step() if i == 0: first_loss = loss.item() current_lr = scheduler.get_last_lr()[0] progress_bar.set_postfix( loss=f"{first_loss:.2f}→{loss.item():.2f}", grad=f"{latents.grad.norm().item():.2e}", lr=f"{current_lr:.2e}", refine_step=f"{i + 1}/{self.num_refinements}", ) store.reset() return latents.detach() class AttentionStore: def __init__(self): self.accumulator = None self.count = 0 self.keep_heads = False def __call__(self, probs: torch.Tensor) -> torch.Tensor: if probs.shape[0] == 2: probs = probs[1:] if not self.keep_heads: probs = probs.mean(dim=1) if self.accumulator is None: self.accumulator = probs else: self.accumulator = self.accumulator + probs self.count += 1 return probs def reset(self): self.accumulator = None self.count = 0 def get_avg_attention(self) -> torch.Tensor: return self.accumulator / self.count class AttnProcessor: r""" Processor for implementing attention (SDPA is used by default if you're using PyTorch 2.0) for the LTX-2.0 model. Compared to the LTX-1.0 model, we allow the RoPE embeddings for the queries and keys to be separate so that we can support audio-to-video (a2v) and video-to-audio (v2a) cross attention. FIXED: Now uses manual attention output for cross-attention to maintain gradient flow. """ _attention_backend = None _parallel_config = None def __init__(self, store: AttentionStore, name: str): if is_torch_version("<", "2.0"): raise ValueError( "LTX attention processors require a minimum PyTorch version of 2.0. Please upgrade your PyTorch installation." ) self.store = store self.name = name def __call__( self, attn: "LTX2Attention", hidden_states: torch.Tensor, encoder_hidden_states: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, query_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, key_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: batch_size, sequence_length, _ = ( hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape ) if attention_mask is not None: attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) original_encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is None: encoder_hidden_states = hidden_states query = attn.to_q(hidden_states) key = attn.to_k(encoder_hidden_states) value = attn.to_v(encoder_hidden_states) query = attn.norm_q(query) key = attn.norm_k(key) if query_rotary_emb is not None: if attn.rope_type == "interleaved": query = apply_interleaved_rotary_emb(query, query_rotary_emb) key = apply_interleaved_rotary_emb( key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb ) elif attn.rope_type == "split": query = apply_split_rotary_emb(query, query_rotary_emb) key = apply_split_rotary_emb(key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb) query = query.unflatten(2, (attn.heads, -1)) key = key.unflatten(2, (attn.heads, -1)) value = value.unflatten(2, (attn.heads, -1)) is_cross = original_encoder_hidden_states is not None and original_encoder_hidden_states is not hidden_states if is_cross: q = query.permute(0, 2, 1, 3) k = key.permute(0, 2, 1, 3) scale_factor = 1.0 / math.sqrt(q.size(-1)) scores = torch.matmul(q, k.transpose(-1, -2)) * scale_factor if attention_mask is not None: scores = scores + attention_mask probs = F.softmax(scores, dim=-1) self.store(probs) hidden_states = dispatch_attention_fn( query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False, backend=self._attention_backend, ) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) hidden_states = attn.to_out[0](hidden_states) hidden_states = attn.to_out[1](hidden_states) return hidden_states 

Hoped someone can point me to my issue, I have a feeling the the captured attention are not affecting the latents during the forward (graph computation).