[codex] Add Mini Studio web UI and fix LTX config loading

.gitignore

@@ -8,3 +8,4 @@ config.json
 **.pyc
 __pycache__/*
 *.egg-info/*
+outputs/

docs/WEBUI-GUIDE.md

@@ -0,0 +1,108 @@
+# mlx-video Mini Studio Guide
+
+Mini Studio is a local Gradio UI for running `mlx-video` generation jobs from a browser while keeping persistent job history, logs, and generated videos.
+
+## Start
+
+```bash
+cd /path/to/mlx-video
+source .venv/bin/activate
+mlx_video.webui --server-port 7861
+```
+
+Open:
+
+```text
+http://127.0.0.1:7861
+```
+
+## Where State Is Stored
+
+Mini Studio writes runtime state under `outputs/`:
+
+- `outputs/jobs.jsonl` — persistent job history.
+- `outputs/logs/<job-id>.log` — stdout/stderr for each generation command.
+- `outputs/videos/<job-id>.mp4` — generated videos.
+
+`outputs/` is ignored by git because it contains local runtime artifacts and large media files.
+
+## Basic LTX-2 Flow
+
+1. Open `Generate -> LTX-2`.
+2. Fill `Prompt`.
+3. Optionally upload an image for image-to-video.
+4. Start with conservative settings:
+   - `Pipeline`: `distilled`
+   - `Width`: `512`
+   - `Height`: `512`
+   - `Frames`: `97`
+   - `Steps`: `30`
+5. Click `Add LTX job to persistent queue` once.
+6. Open `Queue` and click `Refresh`.
+7. Watch `status` and `Latest log tail`.
+8. When the job is `done`, open `Gallery` and click `Refresh gallery`.
+
+## Status Meanings
+
+- `queued` — the job was created and is waiting.
+- `running` — a subprocess is currently executing.
+- `done` — the command exited successfully and should have produced an MP4.
+- `failed` — the command exited with an error; inspect the job log.
+
+## How To Understand What Is Happening
+
+The most useful signal is the log file for the current job:
+
+```bash
+tail -f outputs/logs/<job-id>.log
+```
+
+Typical stages:
+
+1. Model download from Hugging Face.
+2. Text encoder loading.
+3. Transformer/model loading.
+4. Denoising/generation.
+5. Video writing.
+
+The first run can spend a long time downloading model weights. If Hugging Face warns about unauthenticated requests, setting `HF_TOKEN` can improve rate limits.
+
+## Wan 2.x Notes
+
+Wan generation requires an already converted MLX model directory. Mini Studio does not download or convert Wan weights in the first version.
+
+Use the `Models / Settings` tab to check whether a local Wan model directory exists.
+
+## Troubleshooting
+
+If a job appears stuck:
+
+1. Click `Queue -> Refresh`.
+2. Check the newest job status.
+3. Read the log:
+
+   ```bash
+   tail -n 120 outputs/logs/<job-id>.log
+   ```
+
+4. Confirm whether a generation process is alive:
+
+   ```bash
+   ps -axo pid,ppid,stat,etime,%cpu,%mem,command | rg "mlx_video|ltx_2|wan_2|webui"
+   ```
+
+If a job fails after downloading LTX-2 with:
+
+```text
+TypeError: LTXModelConfig.__init__() got an unexpected keyword argument '_class_name'
+```
+
+the local fix is to load LTX config through `LTXModelConfig.from_dict()`, which ignores Hugging Face metadata keys that are not dataclass fields.
+
+If the next failure is:
+
+```text
+ValueError: Received 1680 parameters not in model
+```
+
+the transformer checkpoint keys are still in Diffusers naming form. The local fix is to run unprefixed Hugging Face transformer weights through the same key sanitizer used for raw prefixed checkpoints. The sanitizer normalizes names such as `linear_1` to `linear1`, `to_out.0` to `to_out`, `norm_q` to `q_norm`, and maps top-level Diffusers modules like `time_embed` and `proj_in` to the local MLX module names.

mlx_video/__init__.py

@@ -1,50 +1,62 @@
-from mlx_video.models.ltx_2 import LTXModel, LTXModelConfig
-
-# Audio VAE components
-from mlx_video.models.ltx_2.audio_vae import (
-    AudioDecoder,
-    AudioEncoder,
-    AudioLatentShape,
-    AudioPatchifier,
-    PerChannelStatistics,
-    Vocoder,
-    decode_audio,
-)
-
-# Conditioning
-from mlx_video.models.ltx_2.conditioning import VideoConditionByLatentIndex
-
-# Utilities
-from mlx_video.models.ltx_2.utils import (
-    convert_audio_encoder,
-    get_model_path,
-    load_config,
-    load_safetensors,
-    save_weights,
-)
-from mlx_video.models.wan_2 import WanModel, WanModelConfig
-
 __all__ = [
-    # Models
     "LTXModel",
     "LTXModelConfig",
-    # Audio VAE
     "AudioDecoder",
     "AudioEncoder",
     "Vocoder",
     "decode_audio",
     "AudioPatchifier",
     "AudioLatentShape",
     "PerChannelStatistics",
-    # Conditioning
     "VideoConditionByLatentIndex",
-    # Utilities
     "convert_audio_encoder",
     "get_model_path",
     "load_safetensors",
     "load_config",
     "save_weights",
-    # Wan Models
     "WanModel",
     "WanModelConfig",
 ]
+
+
+def __getattr__(name):
+    if name in {"LTXModel", "LTXModelConfig"}:
+        from mlx_video.models.ltx_2 import LTXModel, LTXModelConfig
+
+        return {"LTXModel": LTXModel, "LTXModelConfig": LTXModelConfig}[name]
+
+    if name in {
+        "AudioDecoder",
+        "AudioEncoder",
+        "AudioLatentShape",
+        "AudioPatchifier",
+        "PerChannelStatistics",
+        "Vocoder",
+        "decode_audio",
+    }:
+        from mlx_video.models.ltx_2 import audio_vae
+
+        return getattr(audio_vae, name)
+
+    if name == "VideoConditionByLatentIndex":
+        from mlx_video.models.ltx_2.conditioning import VideoConditionByLatentIndex
+
+        return VideoConditionByLatentIndex
+
+    if name in {
+        "convert_audio_encoder",
+        "get_model_path",
+        "load_config",
+        "load_safetensors",
+        "save_weights",
+    }:
+        from mlx_video.models.ltx_2 import utils
+
+        return getattr(utils, name)
+
+    if name in {"WanModel", "WanModelConfig"}:
+        from mlx_video.models.wan_2 import WanModel, WanModelConfig
+
+        return {"WanModel": WanModel, "WanModelConfig": WanModelConfig}[name]
+
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

mlx_video/models/ltx_2/config.py

@@ -365,6 +365,13 @@ class VideoEncoderModelConfig(BaseModelConfig):
         ]
     )
 
+    @classmethod
+    def from_dict(cls, params: dict[str, Any]) -> "VideoEncoderModelConfig":
+        params = dict(params)
+        if "latent_channels" in params:
+            params["out_channels"] = params["latent_channels"]
+        return super().from_dict(params)
+
     def __post_init__(self):
         from mlx_video.models.ltx_2.video_vae.convolution import PaddingModeType
         from mlx_video.models.ltx_2.video_vae.resnet import NormLayerType

mlx_video/models/ltx_2/ltx_2.py

@@ -632,32 +632,70 @@ def sanitize(self, weights: dict) -> dict:
         sanitized = {}
 
         has_raw_prefix = any(k.startswith("model.diffusion_model.") for k in weights)
-        if not has_raw_prefix:
-            return weights
+        key_replacements = [
+            (".to_out.0", ".to_out"),
+            (".ff.net.0.proj", ".ff.proj_in"),
+            (".ff.net.2", ".ff.proj_out"),
+            (".audio_ff.net.0.proj", ".audio_ff.proj_in"),
+            (".audio_ff.net.2", ".audio_ff.proj_out"),
+            (".linear_1", ".linear1"),
+            (".linear_2", ".linear2"),
+            (".norm_q", ".q_norm"),
+            (".norm_k", ".k_norm"),
+        ]
+        prefix_replacements = [
+            ("audio_proj_in.", "audio_patchify_proj."),
+            ("proj_in.", "patchify_proj."),
+            ("audio_time_embed.", "audio_adaln_single."),
+            ("time_embed.", "adaln_single."),
+            (
+                "av_cross_attn_video_scale_shift.",
+                "av_ca_video_scale_shift_adaln_single.",
+            ),
+            (
+                "av_cross_attn_audio_scale_shift.",
+                "av_ca_audio_scale_shift_adaln_single.",
+            ),
+            ("av_cross_attn_video_a2v_gate.", "av_ca_a2v_gate_adaln_single."),
+            ("av_cross_attn_audio_v2a_gate.", "av_ca_v2a_gate_adaln_single."),
+        ]
+        exact_replacements = [
+            (
+                ".video_a2v_cross_attn_scale_shift_table",
+                ".scale_shift_table_a2v_ca_video",
+            ),
+            (
+                ".audio_a2v_cross_attn_scale_shift_table",
+                ".scale_shift_table_a2v_ca_audio",
+            ),
+        ]
 
         for key, value in weights.items():
             new_key = key
 
-            if not key.startswith("model.diffusion_model."):
-                continue
-            if (
-                "audio_embeddings_connector" in key
-                or "video_embeddings_connector" in key
-            ):
-                continue
-
-            # Remove 'model.diffusion_model.' prefix
-            new_key = new_key.replace("model.diffusion_model.", "")
-
-            new_key = new_key.replace(".to_out.0.", ".to_out.")
-
-            new_key = new_key.replace(".ff.net.0.proj.", ".ff.proj_in.")
-            new_key = new_key.replace(".ff.net.2.", ".ff.proj_out.")
-            new_key = new_key.replace(".audio_ff.net.0.proj.", ".audio_ff.proj_in.")
-            new_key = new_key.replace(".audio_ff.net.2.", ".audio_ff.proj_out.")
-
-            new_key = new_key.replace(".linear_1.", ".linear1.")
-            new_key = new_key.replace(".linear_2.", ".linear2.")
+            if has_raw_prefix:
+                if not key.startswith("model.diffusion_model."):
+                    continue
+                if (
+                    "audio_embeddings_connector" in key
+                    or "video_embeddings_connector" in key
+                ):
+                    continue
+
+                # Remove 'model.diffusion_model.' prefix
+                new_key = new_key.replace("model.diffusion_model.", "")
+
+            for old, new in key_replacements:
+                if new_key.endswith(old):
+                    new_key = new_key[: -len(old)] + new
+                else:
+                    new_key = new_key.replace(old + ".", new + ".")
+            for old, new in prefix_replacements:
+                if new_key.startswith(old):
+                    new_key = new + new_key[len(old) :]
+                    break
+            for old, new in exact_replacements:
+                new_key = new_key.replace(old, new)
 
             sanitized[new_key] = value
 
@@ -670,7 +708,7 @@ def from_pretrained(cls, model_path: Path, strict: bool = True) -> "LTXModel":
         config_dict = {}
         with open(model_path / "config.json", "r") as f:
             config_dict = json.load(f)
-        config = LTXModelConfig(**config_dict)
+        config = LTXModelConfig.from_dict(config_dict)
         model = cls(config)
 
         weights = {}