Cleanup binaries before space push

MMaDA/inference/common.py

@@ -57,14 +57,6 @@ def build_uni_prompting(cfg) -> Tuple[UniversalPrompting, AutoTokenizer]:
         cond_dropout_prob=cfg.training.cond_dropout_prob,
         use_reserved_token=True,
     )
-    # Safety: if newer task tokens are missing (e.g., <|ti2ti|>, <|t2ti|>), inject them.
-    for tok in ("<|ti2ti|>", "<|t2ti|>"):
-        if tok not in uni_prompting.sptids_dict:
-            token_id = tokenizer.convert_tokens_to_ids(tok)
-            if token_id is None or token_id == tokenizer.unk_token_id:
-                tokenizer.add_special_tokens({"additional_special_tokens": [tok]})
-                token_id = tokenizer.convert_tokens_to_ids(tok)
-            uni_prompting.sptids_dict[tok] = torch.tensor([token_id])
     return uni_prompting, tokenizer
 
 

MMaDA/inference/gradio_multimodal_demo_inst.py

@@ -1259,79 +1259,79 @@ class OmadaDemo:
             return None, "", f"Failed to encode source image: {exc}"
 
         text_tokens = max(4, min(int(text_tokens), self.max_text_len))
-        prompt_ids = self.uni_prompting.text_tokenizer(instruction_clean)['input_ids']
-        if isinstance(prompt_ids, list) and prompt_ids and isinstance(prompt_ids[0], list):
-            prompt_ids = prompt_ids[0]
-        if len(prompt_ids) == 0 or prompt_ids[0] != self.uni_prompting.text_tokenizer.bos_token_id:
-            prompt_ids = [self.uni_prompting.text_tokenizer.bos_token_id] + prompt_ids
-        prompt_ids = prompt_ids + [self.uni_prompting.text_tokenizer.eos_token_id]
-        prompt_tensor = torch.tensor(prompt_ids, device=self.device, dtype=torch.long)
-
-        def _get_token(key: str):
-            tok = self.uni_prompting.sptids_dict.get(key)
-            if tok is None or tok.numel() == 0:
-                return None
-            return int(tok[0].item())
-
-        ti2ti_id = _get_token('<|ti2ti|>')
-        soi_id = _get_token('<|soi|>')
-        eoi_id = _get_token('<|eoi|>')
-        if ti2ti_id is None or soi_id is None or eoi_id is None:
-            return None, "", "TI2TI special tokens are missing in the tokenizer/config."
-        pad_raw = getattr(self.uni_prompting, "pad_id", 0)
-        pad_id = int(pad_raw if pad_raw is not None else 0)
-
-        img_placeholder = torch.full(
-            (self.image_seq_len,),
+
+        # Build prompts using the same helper as training eval (ti2ti_prompt)
+        placeholder_img = torch.full(
+            (1, self.image_seq_len),
             self.mask_token_id,
             dtype=torch.long,
             device=self.device,
         )
-        text_placeholder = torch.full(
-            (text_tokens,),
-            self.mask_token_id,
-            dtype=torch.long,
-            device=self.device,
+        labels_img_placeholder = torch.full_like(placeholder_img, int(self.uni_prompting.ignore_id))
+        text_mask_bool = torch.ones(text_tokens, device=self.device, dtype=torch.bool)
+
+        input_ids, attention_mask, _ = self.uni_prompting.ti2ti_prompt(
+            prompts=[instruction_clean],
+            source_tokens=src_tokens,
+            masked_target_tokens=placeholder_img,
+            labels_img=labels_img_placeholder,
+            target_texts=[""],
+            target_mask_bools=[text_mask_bool],
+            task_token="<|ti2ti|>",
+        )
+        uncond_ids, uncond_attn, _ = self.uni_prompting.ti2ti_prompt(
+            prompts=[""],
+            source_tokens=src_tokens,
+            masked_target_tokens=placeholder_img,
+            labels_img=labels_img_placeholder,
+            target_texts=[""],
+            target_mask_bools=[text_mask_bool],
+            task_token="<|ti2ti|>",
         )
 
-        src_flat = src_tokens.view(-1)
-        prompt_len = prompt_tensor.numel()
-        img_len = img_placeholder.numel()
-        text_len = text_placeholder.numel()
-
-        prompt_start = 2 + src_flat.numel() + 1
-        prompt_end = prompt_start + prompt_len
-        img_start = prompt_end + 1
-        img_end = img_start + img_len
-        text_start = img_end + 1
-        text_end = text_start + text_len
-
-        seq_parts = [
-            torch.tensor([ti2ti_id, soi_id], device=self.device, dtype=torch.long),
-            src_flat,
-            torch.tensor([eoi_id], device=self.device, dtype=torch.long),
-            prompt_tensor,
-            torch.tensor([soi_id], device=self.device, dtype=torch.long),
-            img_placeholder,
-            torch.tensor([eoi_id], device=self.device, dtype=torch.long),
-            text_placeholder,
-        ]
-        seq = torch.cat(seq_parts, dim=0).unsqueeze(0)
-        attn = torch.ones_like(seq, dtype=torch.long, device=self.device)
-
-        uncond_seq = seq.clone()
-        uncond_attn = attn.clone()
-        uncond_seq[:, prompt_start:prompt_end] = pad_id
-        uncond_attn[:, prompt_start:prompt_end] = 0
+        input_ids = input_ids.to(self.device)
+        attention_mask = attention_mask.to(self.device) if attention_mask is not None else None
+        uncond_ids = uncond_ids.to(self.device)
+        uncond_attn = uncond_attn.to(self.device) if uncond_attn is not None else None
+
+        # Locate spans before generation so we can force attention over desired text length
+        seq_example = input_ids[0]
+        soi_id = int(self.uni_prompting.sptids_dict['<|soi|>'][0].item())
+        eoi_id = int(self.uni_prompting.sptids_dict['<|eoi|>'][0].item())
+        pad_id = int(getattr(self.uni_prompting, "pad_id", 0))
+        text_block_len = text_tokens
+
+        soi_positions = (seq_example == soi_id).nonzero(as_tuple=True)[0]
+        eoi_positions = (seq_example == eoi_id).nonzero(as_tuple=True)[0]
+        img_start = img_end = text_start = None
+        if soi_positions.numel() >= 2:
+            tgt_soi = int(soi_positions[1].item())
+            eoi_after = [int(e.item()) for e in eoi_positions if int(e.item()) > tgt_soi]
+            if eoi_after:
+                tgt_eoi = eoi_after[0]
+                img_start = tgt_soi + 1
+                img_end = min(tgt_eoi, input_ids.shape[1])
+                text_start = tgt_eoi + 1
+        if img_start is None:
+            non_pad = (seq_example != pad_id).nonzero(as_tuple=True)
+            pad_offset = int(non_pad[0][0].item()) if len(non_pad) > 0 and non_pad[0].numel() > 0 else 0
+            img_start = pad_offset + 1 + 1 + self.image_seq_len + 1 + self.uni_prompting.max_text_len + 1
+            img_end = img_start + self.image_seq_len
+            text_start = img_end + 1
+        text_end = min(text_start + text_block_len, input_ids.shape[1])
+        if attention_mask is not None:
+            attention_mask[:, text_start:text_end] = 1
+        if uncond_attn is not None:
+            uncond_attn[:, text_start:text_end] = 1
 
         with torch.no_grad():
             filled_tokens, _ = self.model.ti2ti_generate(
-                input_ids=seq.to(self.device),
-                uncond_input_ids=uncond_seq.to(self.device),
-                attention_mask=attn.to(self.device),
-                uncond_attention_mask=uncond_attn.to(self.device),
+                input_ids=input_ids,
+                uncond_input_ids=uncond_ids,
+                attention_mask=attention_mask,
+                uncond_attention_mask=uncond_attn,
                 temperature=float(temperature),
-                timesteps=int(timesteps_image),
+                timesteps=int(max(timesteps_image, timesteps_text)),
                 timesteps_text=int(timesteps_text),
                 timesteps_image=int(timesteps_image),
                 guidance_scale=float(guidance_scale),
@@ -1346,6 +1346,7 @@ class OmadaDemo:
         if filled_tokens is None:
             return None, "", "TI2TI generation failed."
 
+        # Locate spans like evaluate_ti2ti (target image/text blocks)
         filled_tokens = torch.clamp(
             filled_tokens,
             min=0,
@@ -1358,7 +1359,7 @@ class OmadaDemo:
         except Exception as exc:
             return None, "", f"Failed to decode generated image: {exc}"
 
-        text_slice = slice(text_start, min(text_end, filled_tokens.shape[1]))
+        text_slice = slice(text_start, text_end)
         text_block = filled_tokens[:, text_slice]
         text_vocab = self.text_vocab_size
         mask_id = int(self.mask_token_id)