diff --git a/.gitignore b/.gitignore index 3c2f021..536ef54 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,4 @@ config.json **.pyc __pycache__/* *.egg-info/* +outputs/ diff --git a/docs/WEBUI-GUIDE.md b/docs/WEBUI-GUIDE.md new file mode 100644 index 0000000..5373c34 --- /dev/null +++ b/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/.log` — stdout/stderr for each generation command. +- `outputs/videos/.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/.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/.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. diff --git a/mlx_video/__init__.py b/mlx_video/__init__.py index a04ec64..939767e 100644 --- a/mlx_video/__init__.py +++ b/mlx_video/__init__.py @@ -1,34 +1,6 @@ -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", @@ -36,15 +8,55 @@ "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}") diff --git a/mlx_video/models/ltx_2/config.py b/mlx_video/models/ltx_2/config.py index 3d96ab7..64dc85b 100644 --- a/mlx_video/models/ltx_2/config.py +++ b/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 diff --git a/mlx_video/models/ltx_2/ltx_2.py b/mlx_video/models/ltx_2/ltx_2.py index ec21a6e..f080a9d 100644 --- a/mlx_video/models/ltx_2/ltx_2.py +++ b/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 = {} diff --git a/mlx_video/models/ltx_2/video_vae/decoder.py b/mlx_video/models/ltx_2/video_vae/decoder.py index 7d1b8e3..ab9c589 100644 --- a/mlx_video/models/ltx_2/video_vae/decoder.py +++ b/mlx_video/models/ltx_2/video_vae/decoder.py @@ -365,23 +365,61 @@ def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]: sanitized = {} if "per_channel_statistics.mean" in weights: return weights + saw_decoder_weights = False for key, value in weights.items(): new_key = key - if not key.startswith("vae.") or key.startswith("vae.encoder."): + # Converted checkpoints use vae.decoder.*, while Hugging Face + # unified VAE files use decoder.*. + if key.startswith(("vae.encoder.", "encoder.")): + continue + if not ( + key.startswith("vae.") + or key.startswith("decoder.") + or key.startswith("per_channel_statistics.") + ): continue - if key.startswith("vae.per_channel_statistics."): + if key.startswith(("vae.per_channel_statistics.", "per_channel_statistics.")): # Map per-channel statistics (use exact key matching) - if key == "vae.per_channel_statistics.mean-of-means": + if key in { + "vae.per_channel_statistics.mean-of-means", + "per_channel_statistics.mean-of-means", + }: new_key = "per_channel_statistics.mean" - elif key == "vae.per_channel_statistics.std-of-means": + elif key in { + "vae.per_channel_statistics.std-of-means", + "per_channel_statistics.std-of-means", + }: new_key = "per_channel_statistics.std" else: continue # Skip other statistics keys if key.startswith("vae.decoder."): new_key = key.replace("vae.decoder.", "") + saw_decoder_weights = True + elif key.startswith("decoder."): + new_key = key.replace("decoder.", "") + saw_decoder_weights = True + + parts = new_key.split(".") + if ( + len(parts) >= 3 + and parts[0] == "mid_block" + and parts[1] == "resnets" + ): + new_key = ".".join(["up_blocks", "0", "res_blocks"] + parts[2:]) + elif len(parts) >= 4 and parts[0] == "up_blocks" and parts[1].isdigit(): + block_idx = int(parts[1]) + if parts[2] == "resnets": + new_key = ".".join( + ["up_blocks", str(block_idx * 2 + 2), "res_blocks"] + + parts[3:] + ) + elif parts[2] == "upsamplers" and parts[3] == "0": + new_key = ".".join( + ["up_blocks", str(block_idx * 2 + 1)] + parts[4:] + ) # Handle Conv3d weight transpose: (O, I, D, H, W) -> (O, D, H, W, I) if ".conv.weight" in key and value.ndim == 5: @@ -400,6 +438,13 @@ def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]: new_key = new_key.replace(".conv.bias", ".conv.conv.bias") sanitized[new_key] = value + if ( + saw_decoder_weights + and "per_channel_statistics.mean" not in sanitized + and "per_channel_statistics.std" not in sanitized + ): + sanitized["per_channel_statistics.mean"] = self.per_channel_statistics.mean + sanitized["per_channel_statistics.std"] = self.per_channel_statistics.std return sanitized @classmethod @@ -420,6 +465,16 @@ def from_pretrained( model_path = Path(model_path) config_dict = {} + weight_files = sorted(model_path.glob("*.safetensors")) + if ( + not weight_files + and model_path.name == "decoder" + and model_path.parent.exists() + ): + parent_weight_files = sorted(model_path.parent.glob("*.safetensors")) + if parent_weight_files: + model_path = model_path.parent + weight_files = parent_weight_files # Load config from directory config_path = model_path / "config.json" @@ -428,7 +483,6 @@ def from_pretrained( config_dict = json.load(f) # Load weights from directory - weight_files = sorted(model_path.glob("*.safetensors")) if not weight_files: raise FileNotFoundError(f"No safetensors files found in {model_path}") weights = {} @@ -439,7 +493,10 @@ def from_pretrained( decoder_blocks = cls._infer_blocks(weights) # Determine spatial padding mode from config - spatial_padding_mode_str = config_dict.get("spatial_padding_mode", "reflect") + spatial_padding_mode_str = config_dict.get( + "spatial_padding_mode", + config_dict.get("decoder_spatial_padding_mode", "reflect"), + ) spatial_padding_mode = PaddingModeType(spatial_padding_mode_str) model = cls( diff --git a/mlx_video/models/ltx_2/video_vae/video_vae.py b/mlx_video/models/ltx_2/video_vae/video_vae.py index bd85086..e3a6954 100644 --- a/mlx_video/models/ltx_2/video_vae/video_vae.py +++ b/mlx_video/models/ltx_2/video_vae/video_vae.py @@ -385,8 +385,9 @@ def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]: if "position_ids" in key: continue - # Only process VAE encoder weights - if not key.startswith("vae."): + # Only process VAE encoder weights. Converted checkpoints use + # vae.encoder.*, while Hugging Face unified VAE files use encoder.*. + if not (key.startswith("vae.") or key.startswith("encoder.")): continue # Handle per-channel statistics @@ -399,9 +400,30 @@ def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]: continue elif key.startswith("vae.encoder."): new_key = key.replace("vae.encoder.", "") + elif key.startswith("encoder."): + new_key = key.replace("encoder.", "") else: continue + parts = new_key.split(".") + if len(parts) >= 4 and parts[0] == "down_blocks" and parts[1].isdigit(): + block_idx = int(parts[1]) + if parts[2] == "resnets": + new_key = ".".join( + ["down_blocks", str(block_idx * 2), "res_blocks"] + + parts[3:] + ) + elif parts[2] == "downsamplers" and parts[3] == "0": + new_key = ".".join( + ["down_blocks", str(block_idx * 2 + 1)] + parts[4:] + ) + elif ( + len(parts) >= 3 + and parts[0] == "mid_block" + and parts[1] == "resnets" + ): + new_key = ".".join(["down_blocks", "8", "res_blocks"] + parts[2:]) + # Conv3d: PyTorch (O, I, D, H, W) -> MLX (O, D, H, W, I) if "conv" in new_key.lower() and "weight" in new_key and value.ndim == 5: value = mx.transpose(value, (0, 2, 3, 4, 1)) @@ -427,17 +449,28 @@ def from_pretrained(cls, model_path: Path) -> "VideoEncoder": from mlx_video.models.ltx_2.config import VideoEncoderModelConfig + model_path = Path(model_path) + weight_files = sorted(model_path.glob("*.safetensors")) + if ( + not weight_files + and model_path.name == "encoder" + and model_path.parent.exists() + ): + parent_weight_files = sorted(model_path.parent.glob("*.safetensors")) + if parent_weight_files: + model_path = model_path.parent + weight_files = parent_weight_files + # Load config config_path = model_path / "config.json" if config_path.exists(): with open(config_path) as f: config_dict = json.load(f) - config = VideoEncoderModelConfig(**config_dict) + config = VideoEncoderModelConfig.from_dict(config_dict) else: config = VideoEncoderModelConfig() # Load weights - weight_files = sorted(model_path.glob("*.safetensors")) if not weight_files: if model_path.is_file(): weights = mx.load(str(model_path)) diff --git a/mlx_video/webui/__init__.py b/mlx_video/webui/__init__.py new file mode 100644 index 0000000..20a352a --- /dev/null +++ b/mlx_video/webui/__init__.py @@ -0,0 +1 @@ +"""Local Gradio Web UI helpers for mlx-video.""" diff --git a/mlx_video/webui/app.py b/mlx_video/webui/app.py new file mode 100644 index 0000000..17030be --- /dev/null +++ b/mlx_video/webui/app.py @@ -0,0 +1,306 @@ +from __future__ import annotations + +import argparse +from pathlib import Path +from typing import Any + +import gradio as gr + +from mlx_video.webui.job_store import build_job_paths, latest_jobs, read_log_tail +from mlx_video.webui.runner import GenerationRequest, JobManager + + +DEFAULT_OUTPUT_ROOT = Path("outputs") +MANAGER = JobManager(DEFAULT_OUTPUT_ROOT) + + +def _path_or_none(value: Any) -> Path | None: + if not value: + return None + if isinstance(value, str): + return Path(value) + name = getattr(value, "name", None) + return Path(name) if name else None + + +def _jobs_table() -> list[list[str]]: + rows: list[list[str]] = [] + for job in latest_jobs(MANAGER.store.read_all()): + rows.append( + [ + job.get("id", ""), + job.get("status", ""), + job.get("engine", ""), + job.get("prompt", "")[:96], + job.get("output_path", ""), + job.get("updated_at", ""), + ] + ) + return rows + + +def _latest_video_path() -> str | None: + for job in latest_jobs(MANAGER.store.read_all()): + output = Path(job.get("output_path", "")) + if job.get("status") == "done" and output.exists(): + return str(output) + return None + + +def _latest_log_tail() -> str: + jobs = latest_jobs(MANAGER.store.read_all()) + if not jobs: + return "" + return read_log_tail(jobs[0].get("log_path", ""), lines=120) + + +def enqueue_ltx( + prompt: str, + negative_prompt: str, + pipeline: str, + model_repo: str, + image_file: Any, + audio_file: Any, + width: int, + height: int, + num_frames: int, + seed: int, + fps: int, + steps: int, + cfg_scale: float, +) -> tuple[str, list[list[str]], str | None, str]: + paths = build_job_paths(DEFAULT_OUTPUT_ROOT, "preview") + request = GenerationRequest( + engine="ltx", + prompt=prompt, + negative_prompt=negative_prompt or "", + width=int(width), + height=int(height), + num_frames=int(num_frames), + seed=int(seed), + output_path=paths.video, + ltx_pipeline=pipeline, + ltx_model_repo=model_repo or "Lightricks/LTX-2", + image_path=_path_or_none(image_file), + audio_path=_path_or_none(audio_file), + fps=int(fps), + steps=int(steps) if steps else None, + cfg_scale=float(cfg_scale) if cfg_scale else None, + ) + try: + job = MANAGER.enqueue(request) + except Exception as exc: + return f"Error: {exc}", _jobs_table(), _latest_video_path(), _latest_log_tail() + return f"Queued LTX job {job['id']}", _jobs_table(), _latest_video_path(), _latest_log_tail() + + +def enqueue_wan( + prompt: str, + negative_prompt: str, + model_dir: str, + image_file: Any, + width: int, + height: int, + num_frames: int, + seed: int, + steps: int, + guide_scale: str, + shift: float, + scheduler: str, +) -> tuple[str, list[list[str]], str | None, str]: + paths = build_job_paths(DEFAULT_OUTPUT_ROOT, "preview") + request = GenerationRequest( + engine="wan", + prompt=prompt, + negative_prompt=negative_prompt or "", + width=int(width), + height=int(height), + num_frames=int(num_frames), + seed=int(seed), + output_path=paths.video, + wan_model_dir=Path(model_dir).expanduser() if model_dir else None, + image_path=_path_or_none(image_file), + steps=int(steps) if steps else None, + guide_scale=guide_scale or None, + shift=float(shift) if shift else None, + scheduler=scheduler, + ) + try: + job = MANAGER.enqueue(request) + except Exception as exc: + return f"Error: {exc}", _jobs_table(), _latest_video_path(), _latest_log_tail() + return f"Queued Wan job {job['id']}", _jobs_table(), _latest_video_path(), _latest_log_tail() + + +def refresh_queue() -> tuple[list[list[str]], str | None, str]: + return _jobs_table(), _latest_video_path(), _latest_log_tail() + + +def retry_job(job_id: str) -> tuple[str, list[list[str]], str | None, str]: + try: + job = MANAGER.retry(job_id.strip()) + except Exception as exc: + return f"Error: {exc}", _jobs_table(), _latest_video_path(), _latest_log_tail() + return f"Queued retry job {job['id']}", _jobs_table(), _latest_video_path(), _latest_log_tail() + + +def gallery_choices() -> list[str]: + video_dir = DEFAULT_OUTPUT_ROOT / "videos" + if not video_dir.exists(): + return [] + return [str(path) for path in sorted(video_dir.glob("*.mp4"), key=lambda p: p.stat().st_mtime, reverse=True)] + + +def refresh_gallery() -> tuple[gr.Dropdown, str | None]: + choices = gallery_choices() + return gr.Dropdown(choices=choices, value=choices[0] if choices else None), choices[0] if choices else None + + +def select_gallery_video(path: str | None) -> str | None: + return path or None + + +def check_models(ltx_repo: str, wan_dir: str) -> str: + lines = [ + f"LTX model repo: {ltx_repo or 'Lightricks/LTX-2'}", + "LTX weights will be requested by the CLI/Hugging Face when generation starts.", + ] + if wan_dir: + path = Path(wan_dir).expanduser() + lines.append(f"Wan model dir exists: {path.exists()} ({path})") + else: + lines.append("Wan model dir is empty. Wan generation needs a converted MLX model directory.") + return "\n".join(lines) + + +def build_ui() -> gr.Blocks: + with gr.Blocks(title="mlx-video Mini Studio") as demo: + gr.Markdown("# mlx-video Mini Studio") + gr.Markdown("Local Gradio UI for persistent mlx-video jobs.") + + status = gr.Textbox(label="Status", interactive=False) + latest_video = gr.Video(label="Latest video", interactive=False) + log_tail = gr.Textbox(label="Latest log tail", lines=12, interactive=False) + + with gr.Tab("Generate"): + with gr.Tab("LTX-2"): + ltx_prompt = gr.Textbox(label="Prompt", lines=3) + ltx_negative = gr.Textbox(label="Negative prompt") + with gr.Row(): + ltx_pipeline = gr.Dropdown( + ["distilled", "dev", "dev-two-stage", "dev-two-stage-hq"], + value="distilled", + label="Pipeline", + ) + ltx_model_repo = gr.Textbox(value="Lightricks/LTX-2", label="Model repo") + with gr.Row(): + ltx_width = gr.Number(value=512, label="Width", precision=0) + ltx_height = gr.Number(value=512, label="Height", precision=0) + ltx_frames = gr.Number(value=97, label="Frames", precision=0) + with gr.Row(): + ltx_seed = gr.Number(value=42, label="Seed", precision=0) + ltx_fps = gr.Number(value=24, label="FPS", precision=0) + ltx_steps = gr.Number(value=30, label="Steps", precision=0) + ltx_cfg = gr.Number(value=3.0, label="CFG scale") + with gr.Row(): + ltx_image = gr.File(label="Optional image") + ltx_audio = gr.File(label="Optional audio") + ltx_button = gr.Button("Add LTX job to persistent queue", variant="primary") + + with gr.Tab("Wan 2.x"): + wan_prompt = gr.Textbox(label="Prompt", lines=3) + wan_negative = gr.Textbox(label="Negative prompt") + wan_model_dir = gr.Textbox(label="Converted MLX model directory") + wan_image = gr.File(label="Optional image") + with gr.Row(): + wan_width = gr.Number(value=1280, label="Width", precision=0) + wan_height = gr.Number(value=704, label="Height", precision=0) + wan_frames = gr.Number(value=81, label="Frames", precision=0) + with gr.Row(): + wan_seed = gr.Number(value=-1, label="Seed", precision=0) + wan_steps = gr.Number(value=40, label="Steps", precision=0) + wan_guide = gr.Textbox(value="", label="Guide scale") + wan_shift = gr.Number(value=0, label="Shift") + wan_scheduler = gr.Dropdown(["euler", "dpm++", "unipc"], value="unipc", label="Scheduler") + wan_button = gr.Button("Add Wan job to persistent queue", variant="primary") + + with gr.Tab("Queue"): + queue_table = gr.Dataframe( + headers=["id", "status", "engine", "prompt", "output", "updated"], + value=_jobs_table(), + interactive=False, + label="Persistent jobs", + ) + with gr.Row(): + refresh_button = gr.Button("Refresh") + retry_id = gr.Textbox(label="Job id to retry") + retry_button = gr.Button("Retry job") + + with gr.Tab("Gallery"): + gallery_select = gr.Dropdown(choices=gallery_choices(), label="Generated videos") + gallery_video = gr.Video(label="Selected video", interactive=False) + gallery_refresh = gr.Button("Refresh gallery") + + with gr.Tab("Models / Settings"): + settings_ltx_repo = gr.Textbox(value="Lightricks/LTX-2", label="Default LTX repo") + settings_wan_dir = gr.Textbox(label="Wan model directory to check") + settings_button = gr.Button("Check") + settings_output = gr.Textbox(label="Model hints", lines=8, interactive=False) + + ltx_button.click( + enqueue_ltx, + inputs=[ + ltx_prompt, + ltx_negative, + ltx_pipeline, + ltx_model_repo, + ltx_image, + ltx_audio, + ltx_width, + ltx_height, + ltx_frames, + ltx_seed, + ltx_fps, + ltx_steps, + ltx_cfg, + ], + outputs=[status, queue_table, latest_video, log_tail], + ) + wan_button.click( + enqueue_wan, + inputs=[ + wan_prompt, + wan_negative, + wan_model_dir, + wan_image, + wan_width, + wan_height, + wan_frames, + wan_seed, + wan_steps, + wan_guide, + wan_shift, + wan_scheduler, + ], + outputs=[status, queue_table, latest_video, log_tail], + ) + refresh_button.click(refresh_queue, outputs=[queue_table, latest_video, log_tail]) + retry_button.click(retry_job, inputs=[retry_id], outputs=[status, queue_table, latest_video, log_tail]) + gallery_refresh.click(refresh_gallery, outputs=[gallery_select, gallery_video]) + gallery_select.change(select_gallery_video, inputs=[gallery_select], outputs=[gallery_video]) + settings_button.click(check_models, inputs=[settings_ltx_repo, settings_wan_dir], outputs=[settings_output]) + + return demo + + +def main() -> None: + parser = argparse.ArgumentParser(description="Launch mlx-video Mini Studio") + parser.add_argument("--server-name", default="127.0.0.1") + parser.add_argument("--server-port", type=int, default=7860) + parser.add_argument("--share", action="store_true") + args = parser.parse_args() + build_ui().queue().launch(server_name=args.server_name, server_port=args.server_port, share=args.share) + + +if __name__ == "__main__": + main() diff --git a/mlx_video/webui/job_store.py b/mlx_video/webui/job_store.py new file mode 100644 index 0000000..fd60f6c --- /dev/null +++ b/mlx_video/webui/job_store.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +@dataclass(frozen=True) +class JobPaths: + video: Path + log: Path + + +class JobStore: + def __init__(self, output_root: Path | str = "outputs") -> None: + self.output_root = Path(output_root) + self.jobs_file = self.output_root / "jobs.jsonl" + + def append(self, record: dict[str, Any]) -> None: + self.output_root.mkdir(parents=True, exist_ok=True) + with self.jobs_file.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(record, ensure_ascii=False, sort_keys=True)) + handle.write("\n") + + def read_all(self) -> list[dict[str, Any]]: + if not self.jobs_file.exists(): + return [] + records: list[dict[str, Any]] = [] + with self.jobs_file.open("r", encoding="utf-8") as handle: + for line in handle: + line = line.strip() + if line: + records.append(json.loads(line)) + return records + + +def build_job_paths(output_root: Path | str, job_id: str) -> JobPaths: + root = Path(output_root) + video_dir = root / "videos" + log_dir = root / "logs" + video_dir.mkdir(parents=True, exist_ok=True) + log_dir.mkdir(parents=True, exist_ok=True) + return JobPaths(video=video_dir / f"{job_id}.mp4", log=log_dir / f"{job_id}.log") + + +def latest_jobs(records: list[dict[str, Any]]) -> list[dict[str, Any]]: + latest: dict[str, dict[str, Any]] = {} + for record in records: + latest[record["id"]] = record + return sorted(latest.values(), key=lambda record: record.get("updated_at", ""), reverse=True) + + +def read_log_tail(path: Path | str, lines: int = 80) -> str: + log_path = Path(path) + if not log_path.exists(): + return "" + return "\n".join(log_path.read_text(encoding="utf-8", errors="replace").splitlines()[-lines:]) diff --git a/mlx_video/webui/runner.py b/mlx_video/webui/runner.py new file mode 100644 index 0000000..e15c26b --- /dev/null +++ b/mlx_video/webui/runner.py @@ -0,0 +1,245 @@ +from __future__ import annotations + +import subprocess +import sys +import threading +import uuid +from dataclasses import asdict, dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from mlx_video.webui.job_store import JobStore, build_job_paths, latest_jobs + + +def entrypoint_path(name: str) -> str: + return str(Path(sys.executable).with_name(name)) + + +@dataclass +class GenerationRequest: + engine: str + prompt: str + width: int + height: int + num_frames: int + seed: int + output_path: Path + negative_prompt: str = "" + steps: int | None = None + cfg_scale: float | None = None + image_path: Path | None = None + audio_path: Path | None = None + ltx_pipeline: str = "distilled" + ltx_model_repo: str = "Lightricks/LTX-2" + fps: int = 24 + wan_model_dir: Path | None = None + guide_scale: str | None = None + shift: float | None = None + scheduler: str = "unipc" + + +def utc_now() -> str: + return datetime.now(UTC).isoformat(timespec="seconds").replace("+00:00", "Z") + + +def validate_request(request: GenerationRequest) -> None: + if not request.prompt.strip(): + raise ValueError("Prompt is required") + if request.width <= 0 or request.height <= 0: + raise ValueError("Width and height must be positive") + if request.num_frames <= 0: + raise ValueError("Number of frames must be positive") + if request.image_path and not request.image_path.exists(): + raise ValueError(f"Image file does not exist: {request.image_path}") + if request.audio_path and not request.audio_path.exists(): + raise ValueError(f"Audio file does not exist: {request.audio_path}") + if request.engine == "wan" and (not request.wan_model_dir or not request.wan_model_dir.exists()): + raise ValueError("Wan model directory does not exist") + if request.engine not in {"ltx", "wan"}: + raise ValueError(f"Unsupported engine: {request.engine}") + + +def build_command(request: GenerationRequest) -> list[str]: + validate_request(request) + if request.engine == "ltx": + command = [ + entrypoint_path("mlx_video.ltx_2.generate"), + "--prompt", + request.prompt, + "--pipeline", + request.ltx_pipeline, + "--height", + str(request.height), + "--width", + str(request.width), + "--num-frames", + str(request.num_frames), + "--seed", + str(request.seed), + "--fps", + str(request.fps), + "--model-repo", + request.ltx_model_repo, + ] + if request.negative_prompt: + command.extend(["--negative-prompt", request.negative_prompt]) + if request.steps is not None: + command.extend(["--steps", str(request.steps)]) + if request.cfg_scale is not None: + command.extend(["--cfg-scale", str(request.cfg_scale)]) + if request.image_path: + command.extend(["--image", str(request.image_path)]) + if request.audio_path: + command.extend(["--audio-file", str(request.audio_path)]) + command.extend(["--output-path", str(request.output_path)]) + return command + + command = [ + entrypoint_path("mlx_video.wan_2.generate"), + "--model-dir", + str(request.wan_model_dir), + "--prompt", + request.prompt, + "--width", + str(request.width), + "--height", + str(request.height), + "--num-frames", + str(request.num_frames), + "--seed", + str(request.seed), + "--scheduler", + request.scheduler, + ] + if request.negative_prompt: + command.extend(["--negative-prompt", request.negative_prompt]) + if request.image_path: + command.extend(["--image", str(request.image_path)]) + if request.steps is not None: + command.extend(["--steps", str(request.steps)]) + if request.guide_scale: + command.extend(["--guide-scale", request.guide_scale]) + if request.shift is not None: + command.extend(["--shift", str(request.shift)]) + command.extend(["--output-path", str(request.output_path)]) + return command + + +def request_to_params(request: GenerationRequest) -> dict[str, Any]: + params = asdict(request) + params.pop("output_path", None) + for key, value in list(params.items()): + if isinstance(value, Path): + params[key] = str(value) + return params + + +def request_from_params(params: dict[str, Any], output_path: Path) -> GenerationRequest: + data = dict(params) + for key in ("image_path", "audio_path", "wan_model_dir"): + if data.get(key): + data[key] = Path(data[key]) + data["output_path"] = output_path + return GenerationRequest(**data) + + +def clone_for_retry(job: dict[str, Any], new_id: str, output_root: Path | str) -> dict[str, Any]: + paths = build_job_paths(output_root, new_id) + now = utc_now() + return { + "id": new_id, + "engine": job["engine"], + "prompt": job["prompt"], + "params": job.get("params", {}), + "status": "queued", + "output_path": str(paths.video), + "log_path": str(paths.log), + "created_at": now, + "updated_at": now, + } + + +class JobManager: + def __init__(self, output_root: Path | str = "outputs") -> None: + self.output_root = Path(output_root) + self.store = JobStore(self.output_root) + self._lock = threading.Lock() + self._worker: threading.Thread | None = None + + def enqueue(self, request: GenerationRequest) -> dict[str, Any]: + validate_request(request) + job_id = uuid.uuid4().hex[:12] + paths = build_job_paths(self.output_root, job_id) + request.output_path = paths.video + now = utc_now() + job = { + "id": job_id, + "engine": request.engine, + "prompt": request.prompt, + "params": request_to_params(request), + "status": "queued", + "output_path": str(paths.video), + "log_path": str(paths.log), + "created_at": now, + "updated_at": now, + } + self.store.append(job) + self.start_worker() + return job + + def retry(self, job_id: str) -> dict[str, Any]: + jobs = {job["id"]: job for job in latest_jobs(self.store.read_all())} + if job_id not in jobs: + raise ValueError(f"Unknown job id: {job_id}") + new_job = clone_for_retry(jobs[job_id], uuid.uuid4().hex[:12], self.output_root) + self.store.append(new_job) + self.start_worker() + return new_job + + def start_worker(self) -> None: + with self._lock: + if self._worker and self._worker.is_alive(): + return + self._worker = threading.Thread(target=self._run_pending_jobs, daemon=True) + self._worker.start() + + def _run_pending_jobs(self) -> None: + while True: + job = self._next_pending_job() + if not job: + return + self._run_job(job) + + def _next_pending_job(self) -> dict[str, Any] | None: + for job in reversed(latest_jobs(self.store.read_all())): + if job.get("status") == "queued": + return job + return None + + def _run_job(self, job: dict[str, Any]) -> None: + now = utc_now() + running = dict(job, status="running", updated_at=now) + self.store.append(running) + log_path = Path(job["log_path"]) + request = request_from_params(job["params"], Path(job["output_path"])) + command = build_command(request) + with log_path.open("w", encoding="utf-8") as log: + log.write("$ " + " ".join(command) + "\n\n") + log.flush() + try: + process = subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, text=True) + returncode = process.returncode + except OSError as exc: + log.write(f"\nFailed to start command: {exc}\n") + returncode = 127 + status = "done" if returncode == 0 else "failed" + finished = dict( + running, + status=status, + returncode=returncode, + updated_at=utc_now(), + ) + if status == "failed": + finished["error"] = f"Command exited with status {returncode}" + self.store.append(finished) diff --git a/pyproject.toml b/pyproject.toml index 6a4d3ad..d4e4add 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,7 @@ dependencies = [ "imageio>=2.37.2", "imageio-ffmpeg>=0.6.0", "ftfy", + "gradio>=6.0.0", ] license = {text="MIT"} authors = [ @@ -48,6 +49,7 @@ Issues = "https://github.com/Blaizzy/mlx-video/issues" [project.scripts] "mlx_video.ltx_2.generate" = "mlx_video.models.ltx_2.generate:main" "mlx_video.wan_2.generate" = "mlx_video.models.wan_2.generate:main" +"mlx_video.webui" = "mlx_video.webui.app:main" [tool.setuptools.packages.find] include = ["mlx_video*"] @@ -59,4 +61,3 @@ version = {attr = "mlx_video.version.__version__"} dev = [ "pytest", ] - diff --git a/tests/test_ltx_model_loading.py b/tests/test_ltx_model_loading.py new file mode 100644 index 0000000..7ccb9c5 --- /dev/null +++ b/tests/test_ltx_model_loading.py @@ -0,0 +1,83 @@ +import json + + +def test_from_pretrained_ignores_huggingface_config_metadata(tmp_path, monkeypatch): + from mlx_video.models.ltx_2.config import LTXModelConfig + from mlx_video.models.ltx_2.ltx_2 import LTXModel + import mlx_video.models.ltx_2.ltx_2 as ltx_module + + (tmp_path / "config.json").write_text( + json.dumps( + { + "_class_name": "LTXTransformer3DModel", + "_diffusers_version": "0.35.0", + "model_type": "ltx av model", + } + ), + encoding="utf-8", + ) + + captured = {} + + def fake_init(self, config): + captured["config"] = config + + monkeypatch.setattr(LTXModel, "__init__", fake_init) + monkeypatch.setattr(LTXModel, "sanitize", lambda self, weights: weights) + monkeypatch.setattr(LTXModel, "load_weights", lambda self, weights, strict=True: None) + monkeypatch.setattr(LTXModel, "parameters", lambda self: {}) + monkeypatch.setattr(LTXModel, "eval", lambda self: None) + monkeypatch.setattr(ltx_module.mx, "eval", lambda *args, **kwargs: None) + + LTXModel.from_pretrained(tmp_path, strict=False) + + assert isinstance(captured["config"], LTXModelConfig) + assert captured["config"].model_type.value == "ltx av model" + + +def test_sanitize_normalizes_unprefixed_diffusers_transformer_keys(): + from mlx_video.models.ltx_2.ltx_2 import LTXModel + + weights = { + "caption_projection.linear_1.weight": "caption-linear-1", + "caption_projection.linear_2.bias": "caption-linear-2", + "transformer_blocks.0.attn1.to_out.0.weight": "attention-output", + "transformer_blocks.0.ff.net.0.proj.weight": "video-ff-in", + "transformer_blocks.0.ff.net.2.bias": "video-ff-out", + "transformer_blocks.0.audio_ff.net.0.proj.weight": "audio-ff-in", + "transformer_blocks.0.audio_ff.net.2.bias": "audio-ff-out", + "proj_in.weight": "video-patchify", + "audio_proj_in.bias": "audio-patchify", + "time_embed.linear.weight": "video-adaln", + "audio_time_embed.linear.bias": "audio-adaln", + "av_cross_attn_video_scale_shift.linear.weight": "video-scale-shift", + "av_cross_attn_audio_scale_shift.linear.bias": "audio-scale-shift", + "av_cross_attn_video_a2v_gate.linear.weight": "a2v-gate", + "av_cross_attn_audio_v2a_gate.linear.bias": "v2a-gate", + "transformer_blocks.0.attn1.norm_q.weight": "q-norm", + "transformer_blocks.0.attn1.norm_k.weight": "k-norm", + "transformer_blocks.0.video_a2v_cross_attn_scale_shift_table": "video-ca-table", + "transformer_blocks.0.audio_a2v_cross_attn_scale_shift_table": "audio-ca-table", + } + + assert LTXModel.sanitize(None, weights) == { + "caption_projection.linear1.weight": "caption-linear-1", + "caption_projection.linear2.bias": "caption-linear-2", + "transformer_blocks.0.attn1.to_out.weight": "attention-output", + "transformer_blocks.0.ff.proj_in.weight": "video-ff-in", + "transformer_blocks.0.ff.proj_out.bias": "video-ff-out", + "transformer_blocks.0.audio_ff.proj_in.weight": "audio-ff-in", + "transformer_blocks.0.audio_ff.proj_out.bias": "audio-ff-out", + "patchify_proj.weight": "video-patchify", + "audio_patchify_proj.bias": "audio-patchify", + "adaln_single.linear.weight": "video-adaln", + "audio_adaln_single.linear.bias": "audio-adaln", + "av_ca_video_scale_shift_adaln_single.linear.weight": "video-scale-shift", + "av_ca_audio_scale_shift_adaln_single.linear.bias": "audio-scale-shift", + "av_ca_a2v_gate_adaln_single.linear.weight": "a2v-gate", + "av_ca_v2a_gate_adaln_single.linear.bias": "v2a-gate", + "transformer_blocks.0.attn1.q_norm.weight": "q-norm", + "transformer_blocks.0.attn1.k_norm.weight": "k-norm", + "transformer_blocks.0.scale_shift_table_a2v_ca_video": "video-ca-table", + "transformer_blocks.0.scale_shift_table_a2v_ca_audio": "audio-ca-table", + } diff --git a/tests/test_ltx_video_vae_loading.py b/tests/test_ltx_video_vae_loading.py new file mode 100644 index 0000000..a6e8369 --- /dev/null +++ b/tests/test_ltx_video_vae_loading.py @@ -0,0 +1,161 @@ +import json + + +def test_video_encoder_falls_back_to_unified_vae_parent(tmp_path, monkeypatch): + from mlx_video.models.ltx_2.config import VideoEncoderModelConfig + from mlx_video.models.ltx_2.video_vae.video_vae import VideoEncoder + import mlx_video.models.ltx_2.video_vae.video_vae as video_vae_module + + vae_dir = tmp_path / "vae" + encoder_dir = vae_dir / "encoder" + vae_dir.mkdir(parents=True) + (vae_dir / "config.json").write_text( + json.dumps( + { + "_class_name": "AutoencoderKLLTX2Video", + "encoder_spatial_padding_mode": "zeros", + } + ), + encoding="utf-8", + ) + (vae_dir / "diffusion_pytorch_model.safetensors").write_text("", encoding="utf-8") + + captured = {} + + def fake_init(self, config): + captured["config"] = config + + def fake_load(path): + captured["loaded_path"] = path + return {"vae.encoder.conv_in.conv.weight": "encoder-weight"} + + monkeypatch.setattr(VideoEncoder, "__init__", fake_init) + monkeypatch.setattr(VideoEncoder, "sanitize", lambda self, weights: weights) + monkeypatch.setattr(VideoEncoder, "load_weights", lambda self, weights, strict=False: None) + monkeypatch.setattr(video_vae_module.mx, "load", fake_load) + + VideoEncoder.from_pretrained(encoder_dir) + + assert captured["loaded_path"] == str(vae_dir / "diffusion_pytorch_model.safetensors") + assert isinstance(captured["config"], VideoEncoderModelConfig) + + +def test_video_decoder_falls_back_to_unified_vae_parent(tmp_path, monkeypatch): + from mlx_video.models.ltx_2.video_vae.decoder import LTX2VideoDecoder + import mlx_video.models.ltx_2.video_vae.decoder as decoder_module + + vae_dir = tmp_path / "vae" + decoder_dir = vae_dir / "decoder" + vae_dir.mkdir(parents=True) + (vae_dir / "config.json").write_text( + json.dumps( + { + "_class_name": "AutoencoderKLLTX2Video", + "decoder_spatial_padding_mode": "reflect", + "timestep_conditioning": False, + } + ), + encoding="utf-8", + ) + (vae_dir / "diffusion_pytorch_model.safetensors").write_text("", encoding="utf-8") + + captured = {} + + def fake_init(self, **kwargs): + captured["kwargs"] = kwargs + + def fake_load(path): + captured["loaded_path"] = path + return {"vae.decoder.conv_in.conv.weight": "decoder-weight"} + + monkeypatch.setattr(LTX2VideoDecoder, "__init__", fake_init) + monkeypatch.setattr(LTX2VideoDecoder, "_infer_blocks", staticmethod(lambda weights: [])) + monkeypatch.setattr(LTX2VideoDecoder, "sanitize", lambda self, weights: weights) + monkeypatch.setattr(LTX2VideoDecoder, "load_weights", lambda self, weights, strict=True: None) + monkeypatch.setattr(decoder_module.mx, "load", fake_load) + + LTX2VideoDecoder.from_pretrained(decoder_dir) + + assert captured["loaded_path"] == str(vae_dir / "diffusion_pytorch_model.safetensors") + assert captured["kwargs"]["spatial_padding_mode"].value == "reflect" + + +def test_video_encoder_sanitize_accepts_unprefixed_unified_vae_keys(): + from mlx_video.models.ltx_2.config import VideoEncoderModelConfig + from mlx_video.models.ltx_2.video_vae.video_vae import VideoEncoder + + encoder = VideoEncoder(VideoEncoderModelConfig()) + + sanitized = encoder.sanitize( + { + "encoder.conv_in.conv.bias": "conv-bias", + "encoder.down_blocks.0.resnets.0.conv1.conv.bias": "resnet-bias", + "encoder.down_blocks.0.downsamplers.0.conv.conv.bias": "downsample-bias", + "encoder.mid_block.resnets.1.conv2.conv.bias": "mid-bias", + "decoder.conv_in.conv.bias": "decoder-bias", + } + ) + + assert sanitized == { + "conv_in.conv.bias": "conv-bias", + "down_blocks.0.res_blocks.0.conv1.conv.bias": "resnet-bias", + "down_blocks.1.conv.conv.bias": "downsample-bias", + "down_blocks.8.res_blocks.1.conv2.conv.bias": "mid-bias", + } + + +def test_video_encoder_config_uses_diffusers_latent_channels(): + from mlx_video.models.ltx_2.config import VideoEncoderModelConfig + + config = VideoEncoderModelConfig.from_dict( + { + "_class_name": "AutoencoderKLLTX2Video", + "in_channels": 3, + "out_channels": 3, + "latent_channels": 128, + "patch_size": 4, + "encoder_spatial_padding_mode": "zeros", + } + ) + + assert config.out_channels == 128 + + +def test_video_decoder_sanitize_accepts_unprefixed_unified_vae_keys(): + from mlx_video.models.ltx_2.video_vae.decoder import LTX2VideoDecoder + + decoder = LTX2VideoDecoder(decoder_blocks=[]) + + sanitized = decoder.sanitize( + { + "encoder.conv_in.conv.bias": "encoder-bias", + "decoder.conv_in.conv.bias": "conv-bias", + "decoder.up_blocks.0.resnets.0.conv1.conv.bias": "resnet-bias", + "decoder.mid_block.resnets.0.conv1.conv.bias": "mid-bias", + "decoder.up_blocks.0.upsamplers.0.conv.conv.bias": "upsample-bias", + "per_channel_statistics.mean-of-means": "mean", + } + ) + + assert sanitized == { + "conv_in.conv.conv.bias": "conv-bias", + "up_blocks.2.res_blocks.0.conv1.conv.conv.bias": "resnet-bias", + "up_blocks.0.res_blocks.0.conv1.conv.conv.bias": "mid-bias", + "up_blocks.1.conv.conv.bias": "upsample-bias", + "per_channel_statistics.mean": "mean", + } + + +def test_video_decoder_sanitize_adds_identity_statistics_when_missing(): + import mlx.core as mx + + from mlx_video.models.ltx_2.video_vae.decoder import LTX2VideoDecoder + + decoder = LTX2VideoDecoder(decoder_blocks=[]) + + sanitized = decoder.sanitize({"decoder.conv_in.conv.bias": mx.zeros((1,))}) + + assert sanitized["per_channel_statistics.mean"].shape == (128,) + assert sanitized["per_channel_statistics.std"].shape == (128,) + assert mx.all(sanitized["per_channel_statistics.mean"] == 0).item() + assert mx.all(sanitized["per_channel_statistics.std"] == 1).item() diff --git a/tests/test_webui_job_store.py b/tests/test_webui_job_store.py new file mode 100644 index 0000000..d74e83f --- /dev/null +++ b/tests/test_webui_job_store.py @@ -0,0 +1,47 @@ +from pathlib import Path + +from mlx_video.webui.job_store import ( + JobStore, + build_job_paths, + latest_jobs, + read_log_tail, +) + + +def test_latest_jobs_keeps_latest_record_per_job(tmp_path: Path): + store = JobStore(tmp_path) + store.append({"id": "job-a", "status": "queued", "updated_at": "2026-05-24T10:00:00Z"}) + store.append({"id": "job-a", "status": "running", "updated_at": "2026-05-24T10:01:00Z"}) + store.append({"id": "job-a", "status": "done", "updated_at": "2026-05-24T10:02:00Z"}) + + jobs = latest_jobs(store.read_all()) + + assert len(jobs) == 1 + assert jobs[0]["id"] == "job-a" + assert jobs[0]["status"] == "done" + + +def test_latest_jobs_sorts_newest_first(tmp_path: Path): + store = JobStore(tmp_path) + store.append({"id": "older", "status": "done", "updated_at": "2026-05-24T10:00:00Z"}) + store.append({"id": "newer", "status": "queued", "updated_at": "2026-05-24T10:05:00Z"}) + + jobs = latest_jobs(store.read_all()) + + assert [job["id"] for job in jobs] == ["newer", "older"] + + +def test_build_job_paths_uses_outputs_subdirectories(tmp_path: Path): + paths = build_job_paths(tmp_path, "abc123") + + assert paths.video == tmp_path / "videos" / "abc123.mp4" + assert paths.log == tmp_path / "logs" / "abc123.log" + assert paths.video.parent.exists() + assert paths.log.parent.exists() + + +def test_read_log_tail_returns_last_lines(tmp_path: Path): + log = tmp_path / "job.log" + log.write_text("one\ntwo\nthree\nfour\n", encoding="utf-8") + + assert read_log_tail(log, lines=2) == "three\nfour" diff --git a/tests/test_webui_runner.py b/tests/test_webui_runner.py new file mode 100644 index 0000000..7360eac --- /dev/null +++ b/tests/test_webui_runner.py @@ -0,0 +1,120 @@ +from pathlib import Path + +import pytest + +from mlx_video.webui.runner import ( + GenerationRequest, + build_command, + clone_for_retry, + validate_request, +) + + +def test_build_ltx_command_includes_selected_options(tmp_path: Path): + image = tmp_path / "start.png" + image.write_text("fake", encoding="utf-8") + request = GenerationRequest( + engine="ltx", + prompt="A neon city at night", + width=768, + height=512, + num_frames=97, + seed=123, + output_path=tmp_path / "out.mp4", + ltx_pipeline="dev", + ltx_model_repo="prince-canuma/LTX-2-dev", + image_path=image, + steps=40, + cfg_scale=3.5, + ) + + command = build_command(request) + + assert Path(command[0]).is_absolute() + assert command[0].endswith("mlx_video.ltx_2.generate") + assert command[1] == "--prompt" + assert "--pipeline" in command + assert "dev" in command + assert "--model-repo" in command + assert "prince-canuma/LTX-2-dev" in command + assert "--image" in command + assert str(image) in command + assert command[-2:] == ["--output-path", str(tmp_path / "out.mp4")] + + +def test_build_wan_command_includes_model_dir_and_scheduler(tmp_path: Path): + model_dir = tmp_path / "wan" + model_dir.mkdir() + request = GenerationRequest( + engine="wan", + prompt="Ocean waves", + width=640, + height=480, + num_frames=81, + seed=42, + output_path=tmp_path / "wan.mp4", + wan_model_dir=model_dir, + steps=12, + guide_scale="3.0,4.0", + shift=12.0, + scheduler="unipc", + ) + + command = build_command(request) + + assert Path(command[0]).is_absolute() + assert command[0].endswith("mlx_video.wan_2.generate") + assert command[1] == "--model-dir" + assert str(model_dir) in command + assert "--scheduler" in command + assert "unipc" in command + assert command[-2:] == ["--output-path", str(tmp_path / "wan.mp4")] + + +def test_validate_request_rejects_missing_prompt(tmp_path: Path): + request = GenerationRequest( + engine="ltx", + prompt="", + width=512, + height=512, + num_frames=25, + seed=1, + output_path=tmp_path / "out.mp4", + ) + + with pytest.raises(ValueError, match="Prompt is required"): + validate_request(request) + + +def test_validate_request_rejects_missing_wan_model_dir(tmp_path: Path): + request = GenerationRequest( + engine="wan", + prompt="A cat", + width=512, + height=512, + num_frames=25, + seed=1, + output_path=tmp_path / "out.mp4", + wan_model_dir=tmp_path / "missing", + ) + + with pytest.raises(ValueError, match="Wan model directory does not exist"): + validate_request(request) + + +def test_clone_for_retry_preserves_parameters_with_new_id(tmp_path: Path): + old = { + "id": "old", + "engine": "ltx", + "prompt": "A cat", + "params": {"width": 512, "height": 512}, + } + + cloned = clone_for_retry(old, "new", tmp_path) + + assert cloned["id"] == "new" + assert cloned["engine"] == "ltx" + assert cloned["prompt"] == "A cat" + assert cloned["params"] == {"width": 512, "height": 512} + assert cloned["status"] == "queued" + assert cloned["output_path"] == str(tmp_path / "videos" / "new.mp4")