ti2ti model

MMaDA/inference/gradio_multimodal_demo_inst.py

@@ -1079,8 +1079,10 @@ class OmadaDemo:
         if not prompt or not prompt.strip():
             return None, "Please provide a text prompt."
 
+        image_seq_len = 729
+        
         image_tokens = torch.full(
-            (1, self.image_seq_len),
+            (1, image_seq_len),
             self.mask_token_id,
             dtype=torch.long,
             device=self.device,
@@ -1107,7 +1109,7 @@ class OmadaDemo:
                 temperature=float(temperature),
                 timesteps=int(timesteps),
                 noise_schedule=self.image_noise_schedule,
-                seq_len=self.image_seq_len,
+                seq_len=image_seq_len,
                 mask_token_id=self.mask_token_id,
                 codebook_size=self.codebook_size,
                 uni_prompting=self.uni_prompting,
@@ -2056,7 +2058,7 @@ def build_demo(app: OmadaDemo, share: bool, server_name: str, server_port: Optio
                             label="Sub-mode",
                         )
                         with gr.Accordion("Generation settings", open=True, elem_classes=["omada-advanced"], visible=True) as t2i_settings:
-                            t2i_timesteps = gr.Slider(4, 128, value=32, label="Timesteps", step=2)
+                            t2i_timesteps = gr.Slider(4, 128, value=64, label="Timesteps", step=2)
                             t2i_temperature = gr.Slider(0.0, 2.0, value=1.0, label="Sampling temperature", step=0.05)
                             t2i_guidance = gr.Slider(0.0, 8.0, value=3.5, label="CFG scale", step=0.1)
                         with gr.Accordion("Editing settings", open=True, elem_classes=["omada-advanced"], visible=False) as i2i_settings:

app.py

@@ -123,6 +123,17 @@ def download_style() -> Path:
 
 def download_checkpoint() -> Path:
     """Download checkpoint snapshot and return an `unwrapped_model` directory."""
+    local_override = os.getenv("MODEL_CHECKPOINT_PATH")
+    if local_override:
+        override_path = Path(local_override).expanduser()
+        if override_path.name != "unwrapped_model":
+            nested = override_path / "unwrapped_model"
+            if nested.is_dir():
+                override_path = nested
+        if not override_path.exists():
+            raise FileNotFoundError(f"MODEL_CHECKPOINT_PATH does not exist: {override_path}")
+        return override_path
+
     repo_id = os.getenv("MODEL_REPO_ID", "jaeikkim/AIDAS-Omni-Modal-Diffusion")
     revision = os.getenv("MODEL_REVISION", "main")
     token = os.getenv("HF_TOKEN")
@@ -779,4 +790,4 @@ with gr.Blocks(
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()