Cleanup binaries before space push

MMaDA/inference/common.py

@@ -56,6 +56,14 @@ 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

@@ -1267,9 +1267,17 @@ class OmadaDemo:
         prompt_ids = prompt_ids + [self.uni_prompting.text_tokenizer.eos_token_id]
         prompt_tensor = torch.tensor(prompt_ids, device=self.device, dtype=torch.long)
 
-        ti2ti_id = int(self.uni_prompting.sptids_dict['<|ti2ti|>'][0].item())
-        soi_id = int(self.uni_prompting.sptids_dict['<|soi|>'][0].item())
-        eoi_id = int(self.uni_prompting.sptids_dict['<|eoi|>'][0].item())
+        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)