"""Frozen source-side model contract for the Anima MLX conversion. This module intentionally records only facts confirmed from local artifacts. Fields that require the original PyTorch pipeline, tokenizer, or scheduler are left as ``None`` instead of guessed. """ from __future__ import annotations from dataclasses import dataclass from typing import Optional @dataclass(frozen=True) class SafetensorsArtifact: name: str path: str tensor_count: int dtype: str payload_bytes: int file_size_bytes: int @dataclass(frozen=True) class TextEncoderContract: artifact: SafetensorsArtifact vocab_size: int hidden_size: int layer_count: int mlp_intermediate_size: int q_proj_out_features: int kv_proj_out_features: int o_proj_in_features: int qk_norm_size: int tokenizer_type: Optional[str] = None tokenizer_vocab_path: Optional[str] = None chat_template: Optional[str] = None max_sequence_length: Optional[int] = None source_max_length: Optional[int] = None padding_side: Optional[str] = None truncation_side: Optional[str] = None pad_token_id: Optional[int] = None bos_token_id: Optional[int] = None eos_token_id: Optional[int] = None rope_theta: Optional[float] = None output_hidden_state: Optional[str] = None pooling_or_norm: Optional[str] = None uses_auxiliary_t5_token_ids: bool = False @dataclass(frozen=True) class DiffusionContract: artifact: SafetensorsArtifact dit_block_count: int hidden_size: int attention_head_dim: int estimated_attention_heads: int cross_attention_context_dim: int x_embedder_input_dim: int final_patch_dim: int timestep_embedding_dim: int adaln_hidden_dim: int adaln_output_dim_per_sublayer: int llm_adapter_block_count: int llm_adapter_vocab_size: int llm_adapter_hidden_size: int llm_adapter_mlp_intermediate_size: int latent_token_shape: Optional[str] = None x_embedder_input_semantics: Optional[str] = None input_latent_channels: Optional[int] = None padding_mask_channels: Optional[int] = None output_latent_channels: Optional[int] = None patch_spatial: Optional[int] = None patch_temporal: Optional[int] = None max_img_h: Optional[int] = None max_img_w: Optional[int] = None max_frames: Optional[int] = None pos_emb_cls: Optional[str] = None pos_emb_learnable: Optional[bool] = None pos_emb_interpolation: Optional[str] = None rope_h_extrapolation_ratio: Optional[float] = None rope_w_extrapolation_ratio: Optional[float] = None rope_t_extrapolation_ratio: Optional[float] = None min_fps: Optional[int] = None max_fps: Optional[int] = None timestep_range: Optional[str] = None timestep_embedding_method: Optional[str] = None adaln_split_order: Optional[str] = None final_layer_adaln_split_order: Optional[str] = None prediction_target: Optional[str] = None scheduler_type: Optional[str] = None scheduler_shift: Optional[float] = None scheduler_multiplier: Optional[float] = None cfg_formula: Optional[str] = None @dataclass(frozen=True) class VAEContract: artifact: SafetensorsArtifact encoder_input_channels: int encoder_head_channels: int decoder_input_channels: int decoder_output_channels: int base_channels: int max_channels: int middle_attention_channels: int latent_scaling_factor: Optional[float] = None latent_layout: Optional[str] = None output_layout: Optional[str] = None output_range: Optional[str] = None temporal_downsample_ratio: Optional[int] = None spatial_downsample_ratio: Optional[int] = None norm_type: Optional[str] = None attention_axes: Optional[str] = None source_vae_class: Optional[str] = None @dataclass(frozen=True) class SourceContract: text_encoder: TextEncoderContract diffusion: DiffusionContract vae: VAEContract source_pipeline_path: Optional[str] = None source_commit_or_version: Optional[str] = None scheduler_config_path: Optional[str] = None tokenizer_config_path: Optional[str] = None TEXT_ENCODER_ARTIFACT = SafetensorsArtifact( name="text_encoder", path="split_files/text_encoders/qwen_3_06b_base-mlx.safetensors", tensor_count=310, dtype="BF16", payload_bytes=1_192_099_840, file_size_bytes=1_192_135_180, ) DIFFUSION_ARTIFACT = SafetensorsArtifact( name="diffusion", path="split_files/diffusion_models/anima-base-v1.0-mlx.safetensors", tensor_count=685, dtype="BF16", payload_bytes=4_182_137_856, file_size_bytes=4_182_218_400, ) VAE_ARTIFACT = SafetensorsArtifact( name="vae", path="split_files/vae/qwen_image_vae-mlx.safetensors", tensor_count=108, dtype="BF16", payload_bytes=146_590_360, file_size_bytes=146_603_060, ) DEFAULT_SOURCE_CONTRACT = SourceContract( text_encoder=TextEncoderContract( artifact=TEXT_ENCODER_ARTIFACT, vocab_size=151_936, hidden_size=1_024, layer_count=28, mlp_intermediate_size=3_072, q_proj_out_features=2_048, kv_proj_out_features=1_024, o_proj_in_features=2_048, qk_norm_size=128, tokenizer_type="ComfyUI AnimaTokenizer with Qwen2Tokenizer for qwen3_06b and T5TokenizerFast token IDs for t5xxl", tokenizer_vocab_path="tokenizers/qwen25_tokenizer", max_sequence_length=131_072, source_max_length=99_999_999, pad_token_id=151_643, bos_token_id=151_643, eos_token_id=151_645, rope_theta=1_000_000.0, output_hidden_state="last", pooling_or_norm="ComfyUI Qwen3_06BModel uses layer_norm_hidden_state=False", uses_auxiliary_t5_token_ids=True, ), diffusion=DiffusionContract( artifact=DIFFUSION_ARTIFACT, dit_block_count=28, hidden_size=2_048, attention_head_dim=128, estimated_attention_heads=16, cross_attention_context_dim=1_024, x_embedder_input_dim=68, final_patch_dim=64, timestep_embedding_dim=2_048, adaln_hidden_dim=256, adaln_output_dim_per_sublayer=6_144, llm_adapter_block_count=6, llm_adapter_vocab_size=32_128, llm_adapter_hidden_size=1_024, llm_adapter_mlp_intermediate_size=4_096, latent_token_shape="B,C,T,H,W input -> B,T,H/2,W/2,D embedded patches", x_embedder_input_semantics="(16 latent channels + 1 padding mask channel) * patch_temporal 1 * patch_spatial 2 * patch_spatial 2 = 68", input_latent_channels=16, padding_mask_channels=1, output_latent_channels=16, patch_spatial=2, patch_temporal=1, max_img_h=240, max_img_w=240, max_frames=128, pos_emb_cls="rope3d", pos_emb_learnable=True, pos_emb_interpolation="crop", rope_h_extrapolation_ratio=4.0, rope_w_extrapolation_ratio=4.0, rope_t_extrapolation_ratio=1.0, min_fps=1, max_fps=30, timestep_range="ComfyUI ModelSamplingDiscreteFlow sigma in [0, 1], timestep = sigma * multiplier", timestep_embedding_method="Cosmos Predict2 Timesteps sinusoidal embedding then TimestepEmbedding with AdaLN-LoRA", adaln_split_order="shift, scale, gate", final_layer_adaln_split_order="shift, scale", prediction_target="flow/CONST denoising head as used by ComfyUI ModelType.FLOW", scheduler_type="ComfyUI ModelSamplingDiscreteFlow", scheduler_shift=3.0, scheduler_multiplier=1.0, ), vae=VAEContract( artifact=VAE_ARTIFACT, encoder_input_channels=3, encoder_head_channels=32, decoder_input_channels=16, decoder_output_channels=3, base_channels=96, max_channels=384, middle_attention_channels=384, latent_layout="B,C,T,H,W", output_layout="B,C,T,H,W before ComfyUI postprocess", temporal_downsample_ratio=4, spatial_downsample_ratio=8, norm_type="WanVAE RMS_norm", source_vae_class="comfy.ldm.wan.vae.WanVAE", ), source_pipeline_path="bundled minimal MLX runtime", source_commit_or_version="ComfyUI 25757a53c93281e8e2462ced8795373f09e675bf", scheduler_config_path="anima_mlx/runtime/scheduler.py", tokenizer_config_path="anima_mlx/runtime/tokenizer.py", )