From e20deeb5eb8b40af57cd5e90f456469e4166023f Mon Sep 17 00:00:00 2001 From: Shreejith SG Date: Fri, 6 Feb 2026 23:41:33 +0000 Subject: [PATCH] fix(glm-image): resolve SSIM issues by fixing T5 glyph shapes and KV cache dimensions --- fastvideo/models/dits/glm_image.py | 8 +- .../stages/glm_image_before_denoising.py | 130 ++++++++++++++---- 2 files changed, 108 insertions(+), 30 deletions(-) diff --git a/fastvideo/models/dits/glm_image.py b/fastvideo/models/dits/glm_image.py index 98a773559..17bc222ed 100644 --- a/fastvideo/models/dits/glm_image.py +++ b/fastvideo/models/dits/glm_image.py @@ -42,8 +42,8 @@ class GlmImageLayerKVCache: self.k_cache = k self.v_cache = v else: - self.k_cache = torch.cat([self.k_cache, k], dim=1) - self.v_cache = torch.cat([self.v_cache, v], dim=1) + self.k_cache = torch.cat([self.k_cache, k], dim=2) + self.v_cache = torch.cat([self.v_cache, v], dim=2) def get(self): return self.k_cache, self.v_cache @@ -401,9 +401,9 @@ class GlmImageAttention(nn.Module): kv_cache.store(key, value) elif kv_cache.mode == "read": k_cache, v_cache = kv_cache.get() - key = torch.cat([k_cache, key], dim=1) if k_cache is not None else key + key = torch.cat([k_cache, key], dim=2) if k_cache is not None else key value = ( - torch.cat([v_cache, value], dim=1) if v_cache is not None else value + torch.cat([v_cache, value], dim=2) if v_cache is not None else value ) elif kv_cache.mode == "skip": pass diff --git a/fastvideo/pipelines/stages/glm_image_before_denoising.py b/fastvideo/pipelines/stages/glm_image_before_denoising.py index e26d4728a..e929aaeb2 100644 --- a/fastvideo/pipelines/stages/glm_image_before_denoising.py +++ b/fastvideo/pipelines/stages/glm_image_before_denoising.py @@ -178,37 +178,115 @@ class GlmImageBeforeDenoisingStage(PipelineStage): batch.prior_token_id.shape, dtype=torch.bool, device=device) # Drop all priors - # 3. Encode prompts with T5 - text_inputs = self.tokenizer(batch.prompt, - padding="max_length", - max_length=512, - truncation=True, - return_tensors="pt").to(device) + # 3. Encode prompts with T5 (Glyphs ONLY) + # SGLang logic: Only encode text inside quotes for rendering. + def get_glyph_texts(text): + text = text[0] if isinstance(text, list) else text + ocr_texts = ( + re.findall(r"'([^']*)'", text) + + re.findall(r"“([^“”]*)”", text) + + re.findall(r'"([^"]*)"', text) + + re.findall(r"「([^「」]*)」", text) + ) + return ocr_texts - prompt_embeds = self.text_encoder( + glyph_texts = get_glyph_texts(batch.prompt) + # If no glyphs, use empty string (or pad token) + texts_to_encode = glyph_texts if len(glyph_texts) > 0 else [""] + + # Helper to encode list of texts and get embeddings matching SGLang style + # SGLang pads to max_length among the batch of glyphs, and masks rest. + # But here 'batch.prompt' is 1 string (BS=1 usually). + # We need to return [1, L, D] or [1, max_glyph_len, D]? + # SGLang: glyph_embeds = outputs.last_hidden_state[attention_mask.bool()].unsqueeze(0) + # It flattens valid tokens? + + # Let's inspect SGLang `_get_glyph_embeds` (Step 627, lines 434-473) + # It creates a batch of glyphs. Encodes them. + # Then `outputs.last_hidden_state[attention_mask.bool()]` -> Selects ONLY valid tokens from all glyphs. + # .unsqueeze(0) -> [1, Total_Valid_Tokens, D] + + text_inputs = self.tokenizer( + texts_to_encode, + padding="longest", # Pad to longest in this micro-batch + max_length=512, + truncation=True, + return_tensors="pt" + ).to(device) + + # Add padding for processing if needed (SGLang adds pad token buffering?) + # SGLang line 450: `[pad] * ... + input_ids` (Left/Right pad?) + # It seems extensive padding logic. + # For now, standard encoding + mask selection should approximate it. + + prompt_outputs = self.text_encoder( input_ids=text_inputs.input_ids, - attention_mask=text_inputs.attention_mask).last_hidden_state.to( - dtype) - - # Prepare for CFG - concatenate [neg, pos] embeddings for batched processing + attention_mask=text_inputs.attention_mask + ) + + # Flatten valid tokens from all glyphs into one sequence + # inputs.attention_mask identifies valid tokens (excluding pad) + valid_mask = text_inputs.attention_mask.bool() + prompt_embeds = prompt_outputs.last_hidden_state[valid_mask].unsqueeze(0).to(dtype) + + # SGLang returns [1, Seq, D]. + batch.attention_mask = None # Attention mask handled by packing? + # Actually SGLang `GlmImageAttention` uses `context_len` logic. + # FastVideo `GlmImageTransformerBlock` expects `encoder_hidden_states` [B, L, D]. + # If L varies, it's fine. + + # Handle Negatives (CFG) if batch.do_classifier_free_guidance: - neg_text_inputs = self.tokenizer(batch.negative_prompt or "", - padding="max_length", - max_length=512, - truncation=True, - return_tensors="pt").to(device) + neg_prompt = batch.negative_prompt or "" + neg_glyph_texts = get_glyph_texts(neg_prompt) + neg_texts_to_encode = neg_glyph_texts if len(neg_glyph_texts) > 0 else [""] + + neg_inputs = self.tokenizer( + neg_texts_to_encode, + padding="longest", + max_length=512, + truncation=True, + return_tensors="pt" + ).to(device) + + neg_outputs = self.text_encoder( + input_ids=neg_inputs.input_ids, + attention_mask=neg_inputs.attention_mask + ) + neg_valid_mask = neg_inputs.attention_mask.bool() + neg_prompt_embeds = neg_outputs.last_hidden_state[neg_valid_mask].unsqueeze(0).to(dtype) + + # Pad to match length for concatenation? + # Or use list? FastVideo handles lists? + # `glm_image.py` (Model) handles `encoder_hidden_states` as Tensor. I need to PAD. + + L_pos = prompt_embeds.shape[1] + L_neg = neg_prompt_embeds.shape[1] + max_L = max(L_pos, L_neg) + + # Pad with zeros + if L_pos < max_L: + pad = torch.zeros((1, max_L - L_pos, prompt_embeds.shape[2]), device=device, dtype=dtype) + prompt_embeds = torch.cat([prompt_embeds, pad], dim=1) + if L_neg < max_L: + pad = torch.zeros((1, max_L - L_neg, neg_prompt_embeds.shape[2]), device=device, dtype=dtype) + neg_prompt_embeds = torch.cat([neg_prompt_embeds, pad], dim=1) + + prompt_embeds = torch.cat([neg_prompt_embeds, prompt_embeds], dim=0) + # Attention mask for transformer? + # We should provide explicit mask if we padded. + # SGLang: `text_attn_mask`. + # Here we construct mask: [1, max_L] + att_mask_pos = torch.ones((1, max_L), device=device) + if L_pos < max_L: att_mask_pos[:, L_pos:] = 0 + + att_mask_neg = torch.ones((1, max_L), device=device) + if L_neg < max_L: att_mask_neg[:, L_neg:] = 0 + + attention_mask = torch.cat([att_mask_neg, att_mask_pos], dim=0) - neg_prompt_embeds = self.text_encoder( - input_ids=neg_text_inputs.input_ids, - attention_mask=neg_text_inputs.attention_mask - ).last_hidden_state.to(dtype) - - # Concatenate for CFG batching: [uncond, cond] - shape (2, L, D) - prompt_embeds = torch.cat([neg_prompt_embeds, prompt_embeds]) - attention_mask = torch.cat( - [neg_text_inputs.attention_mask, text_inputs.attention_mask]) else: - attention_mask = text_inputs.attention_mask + attention_mask = torch.ones((1, prompt_embeds.shape[1]), device=device) batch.prompt_embeds = [prompt_embeds] batch.attention_mask = attention_mask