Add SeFi Image ZeroGPU app

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.py[cod]
+.gradio/
+outputs/
+.env
+sefi-cache/

LICENSE-SEFI-INFERENCE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 SeFi-Image Authors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,13 +1,34 @@
 ---
-title: Sefi Image Zerogpu
-emoji: 🦀
-colorFrom: indigo
-colorTo: gray
+title: SeFi Image ZeroGPU
 sdk: gradio
 sdk_version: 6.19.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+python_version: 3.12
+license: cc-by-nc-4.0
+models:
+  - SeFi-Image/SeFi-Image-1B-Base
+  - SeFi-Image/SeFi-Image-2B-Base
+  - SeFi-Image/SeFi-Image-5B-Base
+  - SeFi-Image/SeFi-Image-1B-turbo
+  - SeFi-Image/SeFi-Image-2B-turbo
+  - SeFi-Image/SeFi-Image-5B-turbo
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SeFi Image ZeroGPU
+
+Public Gradio ZeroGPU Space for SeFi-Image 1B, 2B, and 5B Base/Turbo checkpoints.
+
+The model defaults match the SeFi model cards:
+
+| Family | Models | Steps | Guidance |
+| --- | --- | ---: | ---: |
+| Base | 1B, 2B, 5B | 50 | 4.0 |
+| Turbo | 1B, 2B, 5B | 4 | 1.0 |
+
+The SeFi checkpoints are gated under CC BY-NC 4.0. The Space uses an `HF_TOKEN`
+secret to download model files, so the Space owner account must first be approved
+for each selected SeFi model repository.
+
+The vendored `sefi/` inference package comes from
+`jmliu206/SeFi-Image` at commit `2f02744a187639ee41a296f8177cbbe7e5f333f5`
+and is covered by `LICENSE-SEFI-INFERENCE`.

app.py

@@ -0,0 +1,353 @@
+from __future__ import annotations
+
+import gc
+import os
+import random
+import threading
+import traceback
+from dataclasses import dataclass
+
+import gradio as gr
+import spaces
+import torch
+from huggingface_hub import hf_hub_download
+
+from sefi import SEFIInferencePipeline
+
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
+
+CACHE_DIR = os.getenv(
+    "SEFI_CACHE_DIR",
+    "/data/sefi-cache" if os.path.isdir("/data") else "/tmp/sefi-cache",
+)
+
+
+@dataclass(frozen=True)
+class ModelPreset:
+    label: str
+    repo_id: str
+    family: str
+    steps: int
+    guidance: float
+
+
+MODEL_PRESETS: dict[str, ModelPreset] = {
+    "1b-base": ModelPreset(
+        label="SeFi-Image 1B Base",
+        repo_id="SeFi-Image/SeFi-Image-1B-Base",
+        family="base",
+        steps=50,
+        guidance=4.0,
+    ),
+    "2b-base": ModelPreset(
+        label="SeFi-Image 2B Base",
+        repo_id="SeFi-Image/SeFi-Image-2B-Base",
+        family="base",
+        steps=50,
+        guidance=4.0,
+    ),
+    "5b-base": ModelPreset(
+        label="SeFi-Image 5B Base",
+        repo_id="SeFi-Image/SeFi-Image-5B-Base",
+        family="base",
+        steps=50,
+        guidance=4.0,
+    ),
+    "1b-turbo": ModelPreset(
+        label="SeFi-Image 1B Turbo",
+        repo_id="SeFi-Image/SeFi-Image-1B-turbo",
+        family="turbo",
+        steps=4,
+        guidance=1.0,
+    ),
+    "2b-turbo": ModelPreset(
+        label="SeFi-Image 2B Turbo",
+        repo_id="SeFi-Image/SeFi-Image-2B-turbo",
+        family="turbo",
+        steps=4,
+        guidance=1.0,
+    ),
+    "5b-turbo": ModelPreset(
+        label="SeFi-Image 5B Turbo",
+        repo_id="SeFi-Image/SeFi-Image-5B-turbo",
+        family="turbo",
+        steps=4,
+        guidance=1.0,
+    ),
+}
+
+DEFAULT_MODEL = "1b-turbo"
+TURBO_STEPS = {4, 8, 10}
+
+_MODEL_LOCK = threading.Lock()
+_LOADED_MODEL_KEY: str | None = None
+_LOADED_PIPE: SEFIInferencePipeline | None = None
+
+
+def _hf_token() -> str | None:
+    token = os.getenv("HF_TOKEN")
+    return token.strip() if token and token.strip() else None
+
+
+def _model_choices() -> list[tuple[str, str]]:
+    return [(preset.label, key) for key, preset in MODEL_PRESETS.items()]
+
+
+def _torch_cleanup() -> None:
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
+
+def _clear_loaded_model() -> None:
+    global _LOADED_MODEL_KEY, _LOADED_PIPE
+    _LOADED_PIPE = None
+    _LOADED_MODEL_KEY = None
+    _torch_cleanup()
+
+
+def _load_pipe(model_key: str) -> SEFIInferencePipeline:
+    global _LOADED_MODEL_KEY, _LOADED_PIPE
+
+    preset = MODEL_PRESETS[model_key]
+    with _MODEL_LOCK:
+        if _LOADED_PIPE is not None and _LOADED_MODEL_KEY == model_key:
+            return _LOADED_PIPE
+
+        _clear_loaded_model()
+        pipe = SEFIInferencePipeline.from_pretrained(
+            preset.repo_id,
+            cache_dir=CACHE_DIR,
+            device="cuda",
+            dtype="bf16",
+        )
+        _LOADED_MODEL_KEY = model_key
+        _LOADED_PIPE = pipe
+        return pipe
+
+
+def _friendly_error(exc: BaseException, repo_id: str | None = None) -> str:
+    text = str(exc)
+    lowered = text.lower()
+    gated = (
+        "requires approval" in lowered
+        or "gated" in lowered
+        or "401" in lowered
+        or "403" in lowered
+    )
+    if gated:
+        repo_hint = f" for `{repo_id}`" if repo_id else ""
+        return (
+            f"Model access is not approved{repo_hint}. Open the model page while "
+            "logged in as the Space owner, accept the SeFi non-commercial gate, "
+            "and retry. The Space already has `HF_TOKEN` configured as a secret."
+        )
+
+    return f"{type(exc).__name__}: {text}"
+
+
+def model_defaults(model_key: str):
+    preset = MODEL_PRESETS[model_key]
+    return (
+        gr.update(value=preset.steps),
+        gr.update(value=preset.guidance),
+        (
+            f"Selected `{preset.repo_id}`. Defaults: "
+            f"{preset.steps} steps, guidance {preset.guidance}."
+        ),
+    )
+
+
+def check_access(model_key: str) -> str:
+    preset = MODEL_PRESETS[model_key]
+    try:
+        hf_hub_download(
+            repo_id=preset.repo_id,
+            filename="sefi_config.yaml",
+            cache_dir=CACHE_DIR,
+            token=_hf_token(),
+        )
+    except Exception as exc:
+        return _friendly_error(exc, preset.repo_id)
+
+    return f"Access OK for `{preset.repo_id}`."
+
+
+def estimate_duration(
+    model_key: str,
+    prompt: str,
+    steps: int,
+    guidance_scale: float,
+    width: int,
+    height: int,
+    seed: int,
+    randomize_seed: bool,
+) -> int:
+    del prompt, guidance_scale, width, height, seed, randomize_seed
+    scale_seconds = {"1b": 120, "2b": 180, "5b": 240}
+    key_prefix = model_key.split("-", 1)[0]
+    load_budget = scale_seconds.get(key_prefix, 180)
+    step_budget = int(max(steps, 1)) * (10 if key_prefix == "5b" else 7)
+    return min(max(load_budget + step_budget, 120), 900)
+
+
+@spaces.GPU(duration=estimate_duration)
+def generate(
+    model_key: str,
+    prompt: str,
+    steps: int,
+    guidance_scale: float,
+    width: int,
+    height: int,
+    seed: int,
+    randomize_seed: bool,
+):
+    prompt = prompt.strip()
+    if not prompt:
+        return None, "Enter a prompt.", seed
+
+    preset = MODEL_PRESETS[model_key]
+    steps = int(steps)
+    guidance_scale = float(guidance_scale)
+    width = int(width)
+    height = int(height)
+
+    if preset.family == "turbo":
+        if steps not in TURBO_STEPS:
+            return (
+                None,
+                "Turbo checkpoints currently support 4, 8, or 10 denoising steps.",
+                seed,
+            )
+        if guidance_scale != 1.0:
+            return None, "Turbo checkpoints should use guidance 1.0.", seed
+
+    if randomize_seed:
+        seed = random.randint(0, 2**31 - 1)
+
+    try:
+        if torch.cuda.is_available():
+            torch.backends.cuda.matmul.allow_tf32 = True
+
+        pipe = _load_pipe(model_key)
+        images = pipe(
+            prompt,
+            num_inference_steps=steps,
+            guidance_scale=guidance_scale,
+            height=height,
+            width=width,
+            seed=int(seed),
+        )
+    except Exception as exc:
+        traceback.print_exc()
+        return None, _friendly_error(exc, preset.repo_id), seed
+
+    if not images:
+        return None, "Generation finished without an image.", seed
+
+    return (
+        images[0],
+        (
+            f"Generated with `{preset.repo_id}` at {width}x{height}, "
+            f"{steps} steps, guidance {guidance_scale}, seed {seed}."
+        ),
+        seed,
+    )
+
+
+with gr.Blocks(title="SeFi Image ZeroGPU") as demo:
+    gr.Markdown("# SeFi Image ZeroGPU")
+
+    with gr.Row():
+        with gr.Column(scale=1, min_width=320):
+            model = gr.Dropdown(
+                label="Model",
+                choices=_model_choices(),
+                value=DEFAULT_MODEL,
+                interactive=True,
+            )
+            prompt = gr.Textbox(
+                label="Prompt",
+                value="A blue ceramic mug on a white desk.",
+                lines=4,
+                max_lines=8,
+            )
+            with gr.Row():
+                steps = gr.Slider(
+                    minimum=1,
+                    maximum=60,
+                    step=1,
+                    value=MODEL_PRESETS[DEFAULT_MODEL].steps,
+                    label="Steps",
+                )
+                guidance = gr.Slider(
+                    minimum=1.0,
+                    maximum=8.0,
+                    step=0.1,
+                    value=MODEL_PRESETS[DEFAULT_MODEL].guidance,
+                    label="Guidance",
+                )
+            with gr.Row():
+                width = gr.Slider(
+                    minimum=512,
+                    maximum=1536,
+                    step=16,
+                    value=1024,
+                    label="Width",
+                )
+                height = gr.Slider(
+                    minimum=512,
+                    maximum=1536,
+                    step=16,
+                    value=1024,
+                    label="Height",
+                )
+            with gr.Row():
+                seed = gr.Number(
+                    label="Seed",
+                    value=42,
+                    precision=0,
+                    minimum=0,
+                    maximum=2**31 - 1,
+                )
+                randomize_seed = gr.Checkbox(label="Randomize", value=False)
+            with gr.Row():
+                run = gr.Button("Generate", variant="primary")
+                access = gr.Button("Check Access")
+
+        with gr.Column(scale=1, min_width=360):
+            image = gr.Image(label="Image", type="pil", format="png")
+            status = gr.Markdown(
+                (
+                    f"Selected `{MODEL_PRESETS[DEFAULT_MODEL].repo_id}`. Defaults: "
+                    f"{MODEL_PRESETS[DEFAULT_MODEL].steps} steps, "
+                    f"guidance {MODEL_PRESETS[DEFAULT_MODEL].guidance}."
+                )
+            )
+
+    model.change(model_defaults, inputs=model, outputs=[steps, guidance, status])
+    access.click(check_access, inputs=model, outputs=status)
+    run.click(
+        generate,
+        inputs=[
+            model,
+            prompt,
+            steps,
+            guidance,
+            width,
+            height,
+            seed,
+            randomize_seed,
+        ],
+        outputs=[image, status, seed],
+        api_name="generate",
+        concurrency_limit=1,
+    )
+
+demo.queue(default_concurrency_limit=1)
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,12 @@
+gradio==6.19.0
+spaces>=0.50.0
+torch>=2.9.1
+torchvision>=0.24.0
+diffusers>=0.39.0
+transformers>=5.13.0
+accelerate>=1.12.0
+safetensors>=0.7.0
+huggingface_hub[hf_xet]>=1.22.0
+hf-xet>=1.2.0
+omegaconf>=2.3.0
+pillow>=12.0.0

sefi/README.md

@@ -0,0 +1,17 @@
+# SeFi Python Package
+
+Reusable Python package for SeFi-Image inference. See `../README.md` for
+installation, model checkpoints, and generation examples.
+
+The package includes:
+
+- checkpoint-derived model metadata
+- checkpoint staging
+- pipeline wrapper
+- prompt/output helpers
+- command-line interface
+
+Weights and model-specific config are loaded from a local checkpoint artifact or
+Hugging Face repo id passed through `--checkpoint` or
+`SEFIInferencePipeline.from_pretrained(...)`. The artifact root should include
+`sefi_config.yaml`.

sefi/__init__.py

@@ -0,0 +1,10 @@
+"""SEFI text-to-image inference package."""
+
+from .pipeline import SEFIInferencePipeline
+from .registry import ModelSpec, infer_model_spec
+
+__all__ = [
+    "ModelSpec",
+    "SEFIInferencePipeline",
+    "infer_model_spec",
+]

sefi/builder.py

@@ -0,0 +1,301 @@
+"""Inference-only component builder for SEFI models."""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass
+
+import torch
+
+from .modeling import (
+    Flux2SEFITransformer2DModel,
+    Qwen3VLTextEncoder,
+    TextureLatentCodec,
+    build_texture_vae,
+)
+
+
+SEFI_SCALE_PRESETS = {
+    "0p5b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 12,
+        "num_layers": 3,
+        "num_single_layers": 10,
+        "joint_attention_dim": 6144,
+    },
+    "1b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 16,
+        "num_layers": 4,
+        "num_single_layers": 12,
+        "joint_attention_dim": 6144,
+    },
+    "2b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 20,
+        "num_layers": 4,
+        "num_single_layers": 16,
+        "joint_attention_dim": 6144,
+    },
+    "3b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 22,
+        "num_layers": 5,
+        "num_single_layers": 18,
+        "joint_attention_dim": 7680,
+    },
+    "4b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 24,
+        "num_layers": 5,
+        "num_single_layers": 20,
+        "joint_attention_dim": 7680,
+    },
+    "5b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 26,
+        "num_layers": 6,
+        "num_single_layers": 21,
+        "joint_attention_dim": 7680,
+    },
+    "6b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 28,
+        "num_layers": 6,
+        "num_single_layers": 22,
+        "joint_attention_dim": 7680,
+    },
+    "8b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 30,
+        "num_layers": 7,
+        "num_single_layers": 24,
+        "joint_attention_dim": 7680,
+    },
+    "9b": {
+        "attention_head_dim": 128,
+        "num_attention_heads": 32,
+        "num_layers": 8,
+        "num_single_layers": 24,
+        "joint_attention_dim": 12288,
+    },
+}
+
+SEFI_MODEL_NAME_TO_SCALE = {
+    "flux2-klein-base-0p5b-sefi": "0p5b",
+    "flux2-klein-base-1b-sefi": "1b",
+    "flux2-klein-base-2b-sefi": "2b",
+    "flux2-klein-base-3b-sefi": "3b",
+    "flux2-klein-base-4b-sefi": "4b",
+    "flux2-klein-base-5b-sefi": "5b",
+    "flux2-klein-base-6b-sefi": "6b",
+    "flux2-klein-base-8b-sefi": "8b",
+    "flux2-klein-base-9b-sefi": "9b",
+}
+
+QWEN3VL_TEXT_HIDDEN_DIMS = {
+    "qwen3vl_2b": 2048,
+    "qwen3vl_4b": 2560,
+    "qwen3vl_8b": 4096,
+}
+
+@dataclass
+class SEFIComponents:
+    transformer: torch.nn.Module
+    text_encoder: torch.nn.Module
+    texture_codec: torch.nn.Module
+    noise_scheduler: object
+    pipeline_cls: type
+    semantic_channels: int
+    texture_channels: int
+    total_channels: int
+
+
+def _resolve_transformer_scale(config) -> str:
+    model_cfg = config.model
+    scale = str(model_cfg.get("transformer_scale", "")).strip().lower()
+    if scale:
+        if scale not in set(SEFI_SCALE_PRESETS) | {"custom"}:
+            raise ValueError(
+                "model.transformer_scale must be one of "
+                f"{list(SEFI_SCALE_PRESETS) + ['custom']}. Got: {scale}"
+            )
+        return scale
+
+    model_name = str(model_cfg.model_name)
+    try:
+        return SEFI_MODEL_NAME_TO_SCALE[model_name]
+    except KeyError as exc:
+        raise ValueError(
+            f"Unsupported SEFI model.model_name: {model_name}. "
+            f"Expected one of {sorted(SEFI_MODEL_NAME_TO_SCALE)}."
+        ) from exc
+
+
+def _derive_semantic_channels(config) -> int:
+    value = config.model.get("semantic_channels", None)
+    if value is None:
+        raise ValueError("Config requires model.semantic_channels for inference.")
+    return int(value)
+
+
+def _texture_vae_config_path(texture_vae_cfg) -> str:
+    name = str(texture_vae_cfg.get("name", "")).strip().lower()
+    base_path = str(texture_vae_cfg.get("base_path", "")).strip()
+    if not base_path:
+        raise ValueError("model.texture_vae.base_path is required.")
+    if name == "sd1.5":
+        return os.path.join(base_path, "config.json")
+    if name in {"flux1", "flux2"}:
+        return os.path.join(base_path, "vae", "config.json")
+    raise ValueError(
+        f"Unsupported model.texture_vae.name: {name}. "
+        "Expected one of ['sd1.5', 'flux1', 'flux2']."
+    )
+
+
+def _derive_texture_channels(config) -> int:
+    config_path = _texture_vae_config_path(config.model.texture_vae)
+    if not os.path.isfile(config_path):
+        raise FileNotFoundError(f"Texture VAE config not found: {config_path}")
+    with open(config_path, "r", encoding="utf-8") as handle:
+        texture_vae_config = json.load(handle)
+    latent_channels = texture_vae_config.get("latent_channels", None)
+    if latent_channels is None:
+        raise ValueError(f"Texture VAE config must contain latent_channels: {config_path}")
+    return int(latent_channels) * 4
+
+
+def _derive_text_output_dim(config) -> int:
+    text_cfg = config.model.text_encoder
+    model_name = str(text_cfg.model_name)
+    if model_name not in QWEN3VL_TEXT_HIDDEN_DIMS:
+        raise ValueError(
+            f"Unsupported SEFI text_encoder.model_name: {model_name}. "
+            f"Expected one of {sorted(QWEN3VL_TEXT_HIDDEN_DIMS)}."
+        )
+    hidden_layers = tuple(int(x) for x in text_cfg.hidden_layers)
+    return int(QWEN3VL_TEXT_HIDDEN_DIMS[model_name]) * len(hidden_layers)
+
+
+def text_encoder_signature(config) -> tuple:
+    text_cfg = config.model.text_encoder
+    return (
+        str(text_cfg.model_name),
+        str(text_cfg.get("weights_root", "outputs/model_weights")),
+        int(text_cfg.max_length),
+        tuple(int(x) for x in text_cfg.hidden_layers),
+    )
+
+
+def build_transformer_config(config, *, total_channels: int, text_output_dim: int) -> dict:
+    from diffusers import Flux2Transformer2DModel
+
+    model_cfg = config.model
+    transformer_cfg_path = str(model_cfg.assets.transformer_config_path)
+    transformer_cfg = Flux2Transformer2DModel.load_config(
+        transformer_cfg_path,
+        subfolder="transformer",
+        local_files_only=True,
+    )
+    transformer_cfg = dict(transformer_cfg)
+
+    transformer_scale = _resolve_transformer_scale(config)
+    if transformer_scale == "custom":
+        overrides = model_cfg.get("transformer_overrides", {})
+        required_keys = (
+            "attention_head_dim",
+            "num_attention_heads",
+            "num_layers",
+            "num_single_layers",
+            "joint_attention_dim",
+        )
+        missing = [key for key in required_keys if key not in overrides]
+        if missing:
+            raise ValueError(
+                "model.transformer_overrides is missing required keys for custom "
+                f"SEFI model: {missing}"
+            )
+        for key in required_keys:
+            transformer_cfg[key] = int(overrides[key])
+        if "mlp_ratio" in overrides:
+            transformer_cfg["mlp_ratio"] = float(overrides["mlp_ratio"])
+    else:
+        transformer_cfg.update(SEFI_SCALE_PRESETS[transformer_scale])
+
+    joint_attention_dim = int(transformer_cfg["joint_attention_dim"])
+    if joint_attention_dim != int(text_output_dim):
+        raise ValueError(
+            "Text dimension mismatch: "
+            f"text_encoder output_dim={text_output_dim}, "
+            f"transformer joint_attention_dim={joint_attention_dim}."
+        )
+
+    transformer_cfg["in_channels"] = int(total_channels)
+    transformer_cfg["out_channels"] = int(total_channels)
+    transformer_cfg["guidance_embeds"] = False
+    return transformer_cfg
+
+
+def build_lightweight_transformer(config, *, total_channels: int, text_output_dim: int):
+    transformer_cfg = build_transformer_config(
+        config,
+        total_channels=total_channels,
+        text_output_dim=text_output_dim,
+    )
+    return Flux2SEFITransformer2DModel(
+        backbone_config=transformer_cfg,
+        text_input_dim=int(text_output_dim),
+    )
+
+
+def build_components(config, *, component_dtype: torch.dtype) -> SEFIComponents:
+    from diffusers import FlowMatchEulerDiscreteScheduler, Flux2KleinPipeline
+
+    model_cfg = config.model
+
+    texture_vae = build_texture_vae(
+        model_cfg.texture_vae,
+        torch_dtype=component_dtype,
+    )
+    texture_codec = TextureLatentCodec(
+        texture_vae=texture_vae,
+        texture_vae_name=str(model_cfg.texture_vae.name),
+    )
+
+    noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+        str(model_cfg.assets.scheduler_path),
+        subfolder="scheduler",
+        local_files_only=True,
+    )
+
+    semantic_channels = _derive_semantic_channels(config)
+    texture_channels = int(texture_codec.texture_channels)
+    total_channels = int(semantic_channels + texture_channels)
+
+    text_cfg = model_cfg.text_encoder
+    text_encoder = Qwen3VLTextEncoder(
+        model_name=str(text_cfg.model_name),
+        weights_root=str(text_cfg.get("weights_root", "outputs/model_weights")),
+        max_length=int(text_cfg.max_length),
+        hidden_layers=[int(x) for x in text_cfg.hidden_layers],
+        torch_dtype=component_dtype,
+    )
+
+    transformer = build_lightweight_transformer(
+        config,
+        total_channels=total_channels,
+        text_output_dim=int(text_encoder.output_dim),
+    )
+
+    return SEFIComponents(
+        transformer=transformer,
+        text_encoder=text_encoder,
+        texture_codec=texture_codec,
+        noise_scheduler=noise_scheduler,
+        pipeline_cls=Flux2KleinPipeline,
+        semantic_channels=semantic_channels,
+        texture_channels=texture_channels,
+        total_channels=total_channels,
+    )

sefi/checkpoints.py

@@ -0,0 +1,95 @@
+"""Checkpoint staging helpers for SEFI inference."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+
+CONFIG_FILENAMES = ("sefi_config.yaml", "config.yaml")
+
+
+def _download_hf_snapshot(
+    repo_id: str,
+    *,
+    cache_dir: str | os.PathLike[str],
+) -> str:
+    try:
+        from huggingface_hub import snapshot_download
+    except ImportError as exc:
+        raise RuntimeError(
+            "Checkpoint is not a local path. Install huggingface_hub or pass a "
+            "local --checkpoint path."
+        ) from exc
+
+    return snapshot_download(
+        repo_id=repo_id,
+        cache_dir=str(cache_dir),
+        local_files_only=False,
+    )
+
+
+def checkpoint_root(path: str | os.PathLike[str]) -> Path:
+    resolved = Path(path).expanduser()
+    return resolved if resolved.is_dir() else resolved.parent
+
+
+def resolve_config_path(
+    checkpoint_path: str | os.PathLike[str],
+    config_path: str | os.PathLike[str] | None = None,
+) -> str:
+    root = checkpoint_root(checkpoint_path)
+
+    if config_path:
+        candidate = Path(config_path).expanduser()
+        if not candidate.is_absolute():
+            rooted = root / candidate
+            if rooted.is_file():
+                return str(rooted)
+        if candidate.is_file():
+            return str(candidate)
+        raise FileNotFoundError(f"Config file not found: {config_path}")
+
+    for filename in CONFIG_FILENAMES:
+        candidate = root / filename
+        if candidate.is_file():
+            return str(candidate)
+
+    expected = ", ".join(CONFIG_FILENAMES)
+    raise FileNotFoundError(
+        f"SEFI config not found under checkpoint root {root}. "
+        f"Expected one of: {expected}. Use --config to override."
+    )
+
+
+def ensure_local_path(
+    checkpoint: str,
+    *,
+    cache_dir: str | os.PathLike[str],
+) -> str:
+    if not checkpoint:
+        raise ValueError(
+            "No checkpoint was provided. Pass a local path or Hugging Face repo id "
+            "with --checkpoint."
+        )
+
+    path = Path(checkpoint).expanduser()
+    if path.exists():
+        return str(path)
+
+    return _download_hf_snapshot(
+        checkpoint,
+        cache_dir=cache_dir,
+    )
+
+
+def resolve_checkpoint_to_local(
+    *,
+    checkpoint: str,
+    cache_dir: str | os.PathLike[str],
+) -> tuple[str, str]:
+    local_path = ensure_local_path(
+        checkpoint,
+        cache_dir=cache_dir,
+    )
+    return local_path, checkpoint

sefi/cli.py

@@ -0,0 +1,111 @@
+"""Command line entry point for SeFi-Image inference."""
+
+from __future__ import annotations
+
+import argparse
+from dataclasses import asdict
+
+from .distributed import (
+    build_rank_generator,
+    setup_distributed,
+    shard_indices_interleaved,
+    wait_for_everyone,
+)
+from .io import expand_prompts, load_prompts, save_images, write_manifest
+from .pipeline import SEFIInferencePipeline
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--prompt", default="")
+    parser.add_argument("--prompt-file", default="")
+    parser.add_argument("--output-dir", default="outputs/inference")
+    parser.add_argument("--cache-dir", default="outputs/model_weights/sefi_inference")
+    parser.add_argument(
+        "--checkpoint",
+        required=True,
+        help="Local checkpoint path or Hugging Face repo id.",
+    )
+    parser.add_argument(
+        "--config",
+        default="",
+        help="Optional config path. Defaults to sefi_config.yaml under --checkpoint.",
+    )
+    parser.add_argument("--steps", type=int, default=None)
+    parser.add_argument("--guidance-scale", type=float, default=None)
+    parser.add_argument("--height", type=int, default=None)
+    parser.add_argument("--width", type=int, default=None)
+    parser.add_argument("--batch-size", type=int, default=1)
+    parser.add_argument("--num-images-per-prompt", type=int, default=1)
+    parser.add_argument("--seed", type=int, default=20260616)
+    parser.add_argument("--device", default="")
+    parser.add_argument("--dtype", choices=("bf16", "fp32"), default="")
+    parser.add_argument("--delta-t", type=float, default=None)
+    parser.add_argument("--timestep-shift-alpha", type=float, default=None)
+    parser.add_argument("--debug-assert-schedule", action="store_true")
+    parser.add_argument("--autoguidance-config", default="")
+    parser.add_argument("--autoguidance-checkpoint", default="")
+    parser.add_argument("--guidance-interval-sigma-lo", type=float, default=None)
+    parser.add_argument("--guidance-interval-sigma-hi", type=float, default=None)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = _parse_args()
+    prompts = load_prompts(
+        prompt=args.prompt or None,
+        prompt_file=args.prompt_file or None,
+    )
+    items = expand_prompts(prompts, args.num_images_per_prompt)
+
+    rank, world_size, device, is_main, accelerator = setup_distributed()
+    local_indices = shard_indices_interleaved(len(items), rank, world_size)
+    local_items = [items[index] for index in local_indices]
+    local_prompts = [item.prompt for item in local_items]
+
+    pipe = SEFIInferencePipeline.from_pretrained(
+        args.checkpoint,
+        cache_dir=args.cache_dir,
+        config=args.config or None,
+        device=args.device or str(device),
+        dtype=args.dtype or None,
+        delta_t=args.delta_t,
+        timestep_shift_alpha=args.timestep_shift_alpha,
+        debug_assert_schedule=args.debug_assert_schedule,
+        autoguidance_config=args.autoguidance_config or None,
+        autoguidance_checkpoint=args.autoguidance_checkpoint or None,
+        guidance_interval_sigma_lo=args.guidance_interval_sigma_lo,
+        guidance_interval_sigma_hi=args.guidance_interval_sigma_hi,
+    )
+
+    generator = build_rank_generator(device, args.seed, rank)
+    images = pipe(
+        local_prompts,
+        num_inference_steps=args.steps,
+        guidance_scale=args.guidance_scale,
+        height=args.height,
+        width=args.width,
+        batch_size=args.batch_size,
+        generator=generator,
+    )
+    save_images(output_dir=args.output_dir, items=local_items, images=images, rank=rank)
+    wait_for_everyone(accelerator)
+
+    if is_main:
+        write_manifest(
+            args.output_dir,
+            {
+                "model": pipe.spec.name,
+                "model_spec": asdict(pipe.spec),
+                "checkpoint_path": pipe.checkpoint_path,
+                "checkpoint_uri": pipe.checkpoint_uri,
+                "num_prompts": len(prompts),
+                "num_images": len(items),
+                "seed": args.seed,
+                "world_size": world_size,
+            },
+        )
+
+
+if __name__ == "__main__":
+    main()

sefi/config.py

@@ -0,0 +1,52 @@
+"""Config loading helpers for SEFI inference."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from omegaconf import OmegaConf
+
+
+def _resolve_relative_path(value, base_dir: Path) -> str:
+    if value is None:
+        return value
+    raw = str(value).strip()
+    if not raw:
+        return raw
+    path = Path(raw).expanduser()
+    if path.is_absolute():
+        return str(path)
+    return str(base_dir / path)
+
+
+def _patch_path(config, dotted_path: str, base_dir: Path) -> None:
+    parts = dotted_path.split(".")
+    node = config
+    for part in parts[:-1]:
+        if part not in node:
+            return
+        node = node[part]
+    leaf = parts[-1]
+    if leaf in node:
+        node[leaf] = _resolve_relative_path(node[leaf], base_dir)
+
+
+def _resolve_artifact_paths(config, base_dir: Path):
+    for dotted_path in (
+        "model.assets.transformer_config_path",
+        "model.assets.scheduler_path",
+        "model.texture_vae.base_path",
+        "model.text_encoder.weights_root",
+    ):
+        _patch_path(config, dotted_path, base_dir)
+    return config
+
+
+def load_config(config_path: str | Path):
+    path = Path(config_path).expanduser()
+    if not path.is_file():
+        raise FileNotFoundError(f"Config file not found: {path}")
+    config = OmegaConf.load(path)
+    config = _resolve_artifact_paths(config, path.parent)
+    OmegaConf.resolve(config)
+    return config

sefi/distributed.py

@@ -0,0 +1,35 @@
+"""Small distributed helpers for CLI inference."""
+
+from __future__ import annotations
+
+import torch
+
+
+def setup_distributed():
+    try:
+        from accelerate import Accelerator
+    except ModuleNotFoundError:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        return 0, 1, device, True, None
+
+    accelerator = Accelerator()
+    return (
+        int(accelerator.process_index),
+        int(accelerator.num_processes),
+        accelerator.device,
+        bool(accelerator.is_main_process),
+        accelerator,
+    )
+
+
+def wait_for_everyone(accelerator) -> None:
+    if accelerator is not None:
+        accelerator.wait_for_everyone()
+
+
+def shard_indices_interleaved(total: int, rank: int, world_size: int) -> list[int]:
+    return list(range(int(rank), int(total), int(world_size)))
+
+
+def build_rank_generator(device: torch.device, seed: int, rank: int) -> torch.Generator:
+    return torch.Generator(device=str(device)).manual_seed(int(seed) + int(rank))

sefi/io.py

@@ -0,0 +1,85 @@
+"""Prompt and output helpers for SEFI inference."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Iterable
+
+from PIL import Image
+
+
+@dataclass(frozen=True)
+class GenerationItem:
+    index: int
+    prompt_index: int
+    repeat_index: int
+    prompt: str
+
+    @property
+    def file_stem(self) -> str:
+        if self.repeat_index == 0:
+            return f"{self.prompt_index:06d}"
+        return f"{self.prompt_index:06d}_{self.repeat_index:02d}"
+
+
+def load_prompts(*, prompt: str | None, prompt_file: str | None) -> list[str]:
+    prompts: list[str] = []
+    if prompt:
+        prompts.append(prompt)
+    if prompt_file:
+        with open(prompt_file, "r", encoding="utf-8") as handle:
+            prompts.extend(line.strip() for line in handle if line.strip())
+    if not prompts:
+        raise ValueError("Provide --prompt or --prompt-file.")
+    return prompts
+
+
+def expand_prompts(prompts: Iterable[str], num_images_per_prompt: int) -> list[GenerationItem]:
+    if num_images_per_prompt <= 0:
+        raise ValueError("num_images_per_prompt must be > 0.")
+
+    items: list[GenerationItem] = []
+    index = 0
+    for prompt_index, prompt in enumerate(prompts):
+        for repeat_index in range(num_images_per_prompt):
+            items.append(
+                GenerationItem(
+                    index=index,
+                    prompt_index=prompt_index,
+                    repeat_index=repeat_index,
+                    prompt=prompt,
+                )
+            )
+            index += 1
+    return items
+
+
+def save_images(
+    *,
+    output_dir: str | Path,
+    items: list[GenerationItem],
+    images: list[Image.Image],
+    rank: int = 0,
+) -> None:
+    if len(items) != len(images):
+        raise ValueError(f"items/images length mismatch: {len(items)} != {len(images)}")
+
+    out = Path(output_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    metadata_path = out / f"metadata_rank{rank:03d}.jsonl"
+    with metadata_path.open("a", encoding="utf-8") as meta:
+        for item, image in zip(items, images):
+            image_path = out / f"{item.file_stem}.png"
+            image.save(image_path)
+            row = asdict(item)
+            row["image"] = image_path.name
+            meta.write(json.dumps(row, ensure_ascii=False) + "\n")
+
+
+def write_manifest(output_dir: str | Path, payload: dict) -> None:
+    out = Path(output_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    with (out / "inference_manifest.json").open("w", encoding="utf-8") as handle:
+        json.dump(payload, handle, ensure_ascii=False, indent=2, sort_keys=True)

sefi/modeling/__init__.py

@@ -0,0 +1,13 @@
+"""SEFI inference model components."""
+
+from .flux2_sefi_transformer import Flux2SEFITransformer2DModel
+from .qwen3vl_text_encoder import Qwen3VLTextEncoder
+from .texture_latent_codec import TextureLatentCodec
+from .texture_vae_factory import build_texture_vae
+
+__all__ = [
+    "Flux2SEFITransformer2DModel",
+    "Qwen3VLTextEncoder",
+    "TextureLatentCodec",
+    "build_texture_vae",
+]

sefi/modeling/flux2_sefi_transformer.py

@@ -0,0 +1,247 @@
+"""Flux2 SEFI transformer wrapper with explicit dual timestep embedding."""
+
+import inspect
+from typing import Any
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from diffusers import Flux2Transformer2DModel
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+
+try:
+    from diffusers.utils import USE_PEFT_BACKEND, scale_lora_layers, unscale_lora_layers
+except Exception:  # pragma: no cover - compatibility fallback
+    USE_PEFT_BACKEND = False
+
+    def scale_lora_layers(model, scale):
+        del model, scale
+
+    def unscale_lora_layers(model, scale):
+        del model, scale
+
+class SEFIDualTimestepEmbeddings(nn.Module):
+    """SEFI dual timestep embeddings: concat([emb_sem, emb_tex])."""
+
+    def __init__(self, in_channels: int, embedding_dim: int, bias: bool = False):
+        super().__init__()
+        if embedding_dim % 2 != 0:
+            raise ValueError(
+                f"SEFI dual timestep embedding requires even embedding_dim, got {embedding_dim}."
+            )
+
+        half_dim = embedding_dim // 2
+        self.time_proj = Timesteps(
+            num_channels=int(in_channels),
+            flip_sin_to_cos=True,
+            downscale_freq_shift=0,
+        )
+        self.semantic_embedder = TimestepEmbedding(
+            in_channels=int(in_channels),
+            time_embed_dim=half_dim,
+            sample_proj_bias=bias,
+        )
+        self.texture_embedder = TimestepEmbedding(
+            in_channels=int(in_channels),
+            time_embed_dim=half_dim,
+            sample_proj_bias=bias,
+        )
+
+    def forward(self, timestep_sem: Tensor, timestep_tex: Tensor) -> Tensor:
+        sem_proj = self.time_proj(timestep_sem)
+        tex_proj = self.time_proj(timestep_tex)
+        sem_emb = self.semantic_embedder(sem_proj.to(timestep_sem.dtype))
+        tex_emb = self.texture_embedder(tex_proj.to(timestep_tex.dtype))
+        return torch.cat([sem_emb, tex_emb], dim=-1)
+
+
+class Flux2SEFITransformer2DModel(nn.Module):
+    """Flux2 transformer wrapper for SEFI inference."""
+
+    def __init__(
+        self,
+        backbone_config: dict,
+        text_input_dim: int,
+    ):
+        super().__init__()
+
+        self.backbone = Flux2Transformer2DModel.from_config(backbone_config)
+        # SEFI handles semantic/texture timesteps explicitly and does not reuse guidance semantics.
+        self.backbone.time_guidance_embed = nn.Identity()
+        self._double_mod_img_kwarg, self._double_mod_txt_kwarg = (
+            self._resolve_double_stream_modulation_kwargs()
+        )
+        self._single_mod_kwarg = self._resolve_single_stream_modulation_kwarg()
+
+        self.dual_time_embed = SEFIDualTimestepEmbeddings(
+            in_channels=int(self.backbone.config.timestep_guidance_channels),
+            embedding_dim=int(self.backbone.inner_dim),
+            bias=False,
+        )
+
+        expected_text_dim = int(self.backbone.config.joint_attention_dim)
+        if int(text_input_dim) != expected_text_dim:
+            raise ValueError(
+                f"Text embedding dim mismatch: text={text_input_dim}, "
+                f"transformer expects={expected_text_dim}."
+            )
+
+    def _resolve_double_stream_modulation_kwargs(self) -> tuple[str, str]:
+        if not self.backbone.transformer_blocks:
+            raise ValueError("Flux2 backbone must define at least one double-stream block.")
+        params = inspect.signature(
+            self.backbone.transformer_blocks[0].forward
+        ).parameters
+        if "temb_mod_img" in params and "temb_mod_txt" in params:
+            return "temb_mod_img", "temb_mod_txt"
+        if "temb_mod_params_img" in params and "temb_mod_params_txt" in params:
+            return "temb_mod_params_img", "temb_mod_params_txt"
+        raise ValueError(
+            "Unsupported Flux2TransformerBlock.forward signature. "
+            "Expected temb_mod_img/temb_mod_txt or "
+            "temb_mod_params_img/temb_mod_params_txt."
+        )
+
+    def _resolve_single_stream_modulation_kwarg(self) -> str:
+        if not self.backbone.single_transformer_blocks:
+            raise ValueError("Flux2 backbone must define at least one single-stream block.")
+        params = inspect.signature(
+            self.backbone.single_transformer_blocks[0].forward
+        ).parameters
+        if "temb_mod" in params:
+            return "temb_mod"
+        if "temb_mod_params" in params:
+            return "temb_mod_params"
+        raise ValueError(
+            "Unsupported Flux2SingleTransformerBlock.forward signature. "
+            "Expected temb_mod or temb_mod_params."
+        )
+
+    def _format_single_stream_modulation(self, single_stream_mod):
+        if self._single_mod_kwarg != "temb_mod_params":
+            return single_stream_mod
+        if (
+            isinstance(single_stream_mod, tuple)
+            and len(single_stream_mod) == 1
+            and isinstance(single_stream_mod[0], tuple)
+            and len(single_stream_mod[0]) == 3
+        ):
+            return single_stream_mod[0]
+        return single_stream_mod
+
+    def enable_gradient_checkpointing(self):
+        self.backbone.enable_gradient_checkpointing()
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        timestep_sem: Tensor,
+        timestep_tex: Tensor,
+        encoder_hidden_states: Tensor,
+        txt_ids: Tensor,
+        img_ids: Tensor,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+    ) -> Tensor:
+        model_device = hidden_states.device
+        model_dtype = next(self.backbone.parameters()).dtype
+
+        hidden_states = hidden_states.to(device=model_device, dtype=model_dtype)
+
+        encoder_hidden_states = encoder_hidden_states.to(
+            device=model_device, dtype=model_dtype
+        )
+
+        timestep_sem = timestep_sem.to(device=model_device, dtype=model_dtype)
+        timestep_tex = timestep_tex.to(device=model_device, dtype=model_dtype)
+
+        if joint_attention_kwargs is not None:
+            joint_attention_kwargs = joint_attention_kwargs.copy()
+            lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        # 0) LoRA scaling (keep semantics aligned with Flux2 forward).
+        if USE_PEFT_BACKEND:
+            scale_lora_layers(self.backbone, lora_scale)
+
+        num_txt_tokens = encoder_hidden_states.shape[1]
+        # 1) SEFI dual-time embedding + modulation parameters.
+        temb = self.dual_time_embed(timestep_sem * 1000, timestep_tex * 1000)
+
+        double_stream_mod_img = self.backbone.double_stream_modulation_img(temb)
+        double_stream_mod_txt = self.backbone.double_stream_modulation_txt(temb)
+        single_stream_mod = self.backbone.single_stream_modulation(temb)
+        single_stream_block_mod = self._format_single_stream_modulation(single_stream_mod)
+
+        # 2) Input projection for image/text streams.
+        hidden_states = self.backbone.x_embedder(hidden_states)
+        encoder_hidden_states = self.backbone.context_embedder(encoder_hidden_states)
+
+        if img_ids.ndim == 3:
+            img_ids = img_ids[0]
+        if txt_ids.ndim == 3:
+            txt_ids = txt_ids[0]
+
+        image_rotary_emb = self.backbone.pos_embed(img_ids)
+        text_rotary_emb = self.backbone.pos_embed(txt_ids)
+        concat_rotary_emb = (
+            torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0),
+            torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0),
+        )
+
+        # 3) Double-stream transformer blocks.
+        for block in self.backbone.transformer_blocks:
+            if torch.is_grad_enabled() and self.backbone.gradient_checkpointing:
+                encoder_hidden_states, hidden_states = self.backbone._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    double_stream_mod_img,
+                    double_stream_mod_txt,
+                    concat_rotary_emb,
+                    joint_attention_kwargs,
+                )
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    **{
+                        self._double_mod_img_kwarg: double_stream_mod_img,
+                        self._double_mod_txt_kwarg: double_stream_mod_txt,
+                    },
+                    image_rotary_emb=concat_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        # 4) Single-stream transformer blocks.
+        for block in self.backbone.single_transformer_blocks:
+            if torch.is_grad_enabled() and self.backbone.gradient_checkpointing:
+                hidden_states = self.backbone._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    None,
+                    single_stream_block_mod,
+                    concat_rotary_emb,
+                    joint_attention_kwargs,
+                )
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=None,
+                    **{self._single_mod_kwarg: single_stream_block_mod},
+                    image_rotary_emb=concat_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+        # 5) Output layers.
+        hidden_states = hidden_states[:, num_txt_tokens:, ...]
+        hidden_states = self.backbone.norm_out(hidden_states, temb)
+        model_pred = self.backbone.proj_out(hidden_states)
+
+        if USE_PEFT_BACKEND:
+            unscale_lora_layers(self.backbone, lora_scale)
+
+        return model_pred

sefi/modeling/qwen3vl_text_encoder.py

@@ -0,0 +1,147 @@
+"""Qwen3-VL text encoder wrapper for SEFI T2I inference."""
+
+import os
+from typing import Sequence
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
+
+
+QWEN3VL_MODEL_PATHS = {
+    "qwen3vl_2b": "Qwen3-VL-2B-Instruct",
+    "qwen3vl_4b": "Qwen3-VL-4B-Instruct",
+    "qwen3vl_8b": "Qwen3-VL-8B-Instruct",
+}
+
+
+def resolve_qwen3vl_model_path(
+    model_name: str,
+    weights_root: str = "outputs/model_weights",
+) -> str:
+    if model_name not in QWEN3VL_MODEL_PATHS:
+        raise ValueError(
+            f"Unsupported Qwen3-VL model: {model_name}. "
+            f"Supported: {list(QWEN3VL_MODEL_PATHS.keys())}"
+        )
+
+    model_path = os.path.join(weights_root, QWEN3VL_MODEL_PATHS[model_name])
+    if not os.path.exists(model_path):
+        raise FileNotFoundError(
+            f"Qwen3-VL model not found: {model_path}. "
+            "Please download weights to outputs/model_weights first."
+        )
+    return model_path
+
+
+class Qwen3VLTextEncoder(nn.Module):
+    """Text embedding wrapper using Qwen3-VL language model."""
+
+    def __init__(
+        self,
+        model_name: str,
+        weights_root: str = "outputs/model_weights",
+        max_length: int = 512,
+        hidden_layers: Sequence[int] = (9, 18, 27),
+        torch_dtype: torch.dtype = torch.bfloat16,
+    ):
+        super().__init__()
+
+        self.model_name = model_name
+        self.max_length = int(max_length)
+        self.hidden_layers = tuple(int(x) for x in hidden_layers)
+
+        model_path = resolve_qwen3vl_model_path(model_name, weights_root=weights_root)
+        self.processor = AutoProcessor.from_pretrained(model_path, local_files_only=True)
+        self.tokenizer = self.processor.tokenizer
+
+        self.model = Qwen3VLForConditionalGeneration.from_pretrained(
+            model_path,
+            torch_dtype=torch_dtype,
+            local_files_only=True,
+            device_map="cpu",
+        )
+
+        # Keep only text tower to save memory.
+        if hasattr(self.model, "model") and hasattr(self.model.model, "visual"):
+            del self.model.model.visual
+
+        self.model.eval()
+
+        text_hidden_size = int(self.model.config.text_config.hidden_size)
+        self.output_dim = text_hidden_size * len(self.hidden_layers)
+
+    def _build_chat_text(self, caption: str) -> str:
+        messages = [{"role": "user", "content": [{"type": "text", "text": caption}]}]
+        try:
+            return self.processor.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+                enable_thinking=False,
+            )
+        except TypeError:
+            return self.processor.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+            )
+
+    @staticmethod
+    def _prepare_text_ids(x: Tensor, t_coord: Tensor | None = None) -> Tensor:
+        batch, seq_len, _ = x.shape
+        out_ids = []
+
+        for i in range(batch):
+            t = torch.arange(1) if t_coord is None else t_coord[i]
+            h = torch.arange(1)
+            w = torch.arange(1)
+            l = torch.arange(seq_len)
+            coords = torch.cartesian_prod(t, h, w, l)
+            out_ids.append(coords)
+
+        return torch.stack(out_ids)
+
+    @torch.no_grad()
+    def encode(self, captions: list[str], dtype: torch.dtype | None = None) -> tuple[Tensor, Tensor]:
+        device = next(self.model.parameters()).device
+        model_dtype = next(self.model.parameters()).dtype
+        if dtype is None:
+            dtype = model_dtype
+
+        chat_texts = [self._build_chat_text(caption) for caption in captions]
+        tokenized = self.tokenizer(
+            chat_texts,
+            return_tensors="pt",
+            padding="max_length",
+            truncation=True,
+            max_length=self.max_length,
+        )
+        input_ids = tokenized["input_ids"].to(device)
+        attention_mask = tokenized["attention_mask"].to(device)
+
+        output = self.model.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+            use_cache=False,
+            return_dict=True,
+        )
+
+        hidden_states = output.hidden_states
+        max_idx = len(hidden_states) - 1
+        for layer_idx in self.hidden_layers:
+            if layer_idx > max_idx:
+                raise ValueError(
+                    f"Requested hidden layer {layer_idx}, but model only provides up to {max_idx}."
+                )
+
+        stacked = torch.stack([hidden_states[idx] for idx in self.hidden_layers], dim=1)
+        stacked = stacked.to(dtype=dtype)
+
+        batch, num_layers, seq_len, hidden_dim = stacked.shape
+        prompt_embeds = stacked.permute(0, 2, 1, 3).reshape(batch, seq_len, num_layers * hidden_dim)
+        text_ids = self._prepare_text_ids(prompt_embeds).to(device)
+
+        return prompt_embeds, text_ids

sefi/modeling/texture_latent_codec.py

@@ -0,0 +1,103 @@
+"""Texture latent codec for SEFI-T2I."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+class TextureLatentCodec(nn.Module):
+    """Encode/decode and normalize texture latents for SEFI."""
+
+    def __init__(
+        self,
+        texture_vae: nn.Module,
+        texture_vae_name: str,
+    ):
+        super().__init__()
+        self.texture_vae = texture_vae
+        self.texture_vae_name = str(texture_vae_name)
+        self._use_flux2_bn = self.texture_vae_name == "flux2"
+
+        config = getattr(texture_vae, "config", None)
+        latent_channels = getattr(config, "latent_channels", None)
+        if latent_channels is None:
+            raise ValueError(
+                "Texture VAE config must provide latent_channels for channel derivation."
+            )
+        self.latent_channels = int(latent_channels)
+        self.texture_channels = int(self.latent_channels * 4)
+
+        if self._use_flux2_bn:
+            if not hasattr(texture_vae, "bn"):
+                raise ValueError(
+                    f"Texture VAE '{self.texture_vae_name}' requires bn stats but no bn module found."
+                )
+            eps = float(getattr(config, "batch_norm_eps", 1e-6))
+            bn_mean = texture_vae.bn.running_mean.view(1, -1, 1, 1).float()
+            bn_std = torch.sqrt(texture_vae.bn.running_var.view(1, -1, 1, 1).float() + eps)
+            self.register_buffer("texture_bn_mean", bn_mean, persistent=False)
+            self.register_buffer("texture_bn_std", bn_std, persistent=False)
+            self.scaling_factor = None
+            self.shift_factor = None
+        else:
+            scaling_factor = float(getattr(config, "scaling_factor", 1.0))
+            shift_factor = float(getattr(config, "shift_factor", 0.0) or 0.0)
+            if scaling_factor <= 0:
+                raise ValueError(
+                    f"Invalid scaling_factor={scaling_factor} for texture VAE {self.texture_vae_name}."
+                )
+            self.scaling_factor = scaling_factor
+            self.shift_factor = shift_factor
+
+    @property
+    def vae_dtype(self) -> torch.dtype:
+        return next(self.texture_vae.parameters()).dtype
+
+    @torch.no_grad()
+    def _encode_raw(self, images: Tensor) -> Tensor:
+        return self.texture_vae.encode(images.to(dtype=self.vae_dtype)).latent_dist.mode()
+
+    def _normalize_raw(self, raw_latents: Tensor) -> Tensor:
+        return (raw_latents - self.shift_factor) * self.scaling_factor
+
+    def _denormalize_raw(self, normed_latents: Tensor) -> Tensor:
+        return normed_latents / self.scaling_factor + self.shift_factor
+
+    def _normalize_patchified(self, patchified_latents: Tensor) -> Tensor:
+        bn_mean = self.texture_bn_mean.to(patchified_latents.device, patchified_latents.dtype)
+        bn_std = self.texture_bn_std.to(patchified_latents.device, patchified_latents.dtype)
+        return (patchified_latents - bn_mean) / bn_std
+
+    def _denormalize_patchified(self, patchified_latents: Tensor) -> Tensor:
+        bn_mean = self.texture_bn_mean.to(patchified_latents.device, patchified_latents.dtype)
+        bn_std = self.texture_bn_std.to(patchified_latents.device, patchified_latents.dtype)
+        return patchified_latents * bn_std + bn_mean
+
+    @torch.no_grad()
+    def encode_texture(self, images: Tensor, pipeline_cls) -> Tensor:
+        raw_latents = self._encode_raw(images)
+        if self._use_flux2_bn:
+            patchified = pipeline_cls._patchify_latents(raw_latents)
+            patchified = self._normalize_patchified(patchified)
+        else:
+            normed_raw = self._normalize_raw(raw_latents)
+            patchified = pipeline_cls._patchify_latents(normed_raw)
+
+        if patchified.shape[1] != self.texture_channels:
+            raise ValueError(
+                f"Texture channels mismatch: derived={self.texture_channels}, got={patchified.shape[1]}."
+            )
+
+        return patchified
+
+    @torch.no_grad()
+    def decode_texture(self, texture_latents: Tensor, pipeline_cls) -> Tensor:
+        if self._use_flux2_bn:
+            patchified = self._denormalize_patchified(texture_latents)
+            raw_latents = pipeline_cls._unpatchify_latents(patchified)
+        else:
+            raw_normed = pipeline_cls._unpatchify_latents(texture_latents)
+            raw_latents = self._denormalize_raw(raw_normed)
+        return self.texture_vae.decode(raw_latents, return_dict=False)[0]

sefi/modeling/texture_vae_factory.py

@@ -0,0 +1,41 @@
+"""Texture VAE factory for SEFI-T2I."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+import torch
+
+from .vae_registry import load_vae_from_path
+
+SUPPORTED_TEXTURE_VAE_NAMES = {
+    "sd1.5",
+    "flux1",
+    "flux2",
+}
+
+
+def _normalize_texture_vae_name(name: str) -> str:
+    normalized = str(name).strip().lower()
+    if normalized not in SUPPORTED_TEXTURE_VAE_NAMES:
+        raise ValueError(
+            f"Unsupported model.texture_vae.name={name}. "
+            f"Expected one of {sorted(SUPPORTED_TEXTURE_VAE_NAMES)}."
+        )
+    return normalized
+
+
+def build_texture_vae(texture_vae_cfg: Mapping, *, torch_dtype: torch.dtype):
+    """Build the final texture VAE packaged in a SEFI inference artifact."""
+    name = _normalize_texture_vae_name(str(texture_vae_cfg.get("name", "")))
+    base_path = str(texture_vae_cfg.get("base_path", "")).strip()
+    if not base_path:
+        raise ValueError("model.texture_vae.base_path is required.")
+
+    load_name = "flux2" if name == "flux2" else name
+    return load_vae_from_path(
+        load_name,
+        base_path,
+        torch_dtype=torch_dtype,
+        local_files_only=True,
+    )

sefi/modeling/vae_registry.py

@@ -0,0 +1,42 @@
+"""VAE loading helpers for SEFI inference artifacts."""
+
+from __future__ import annotations
+
+import torch
+
+
+def load_vae_from_path(
+    model_name: str,
+    base_path: str,
+    *,
+    torch_dtype: torch.dtype | None = None,
+    local_files_only: bool = True,
+):
+    """Load a final inference VAE from an explicit artifact path."""
+    import diffusers
+
+    kwargs = {"local_files_only": local_files_only}
+    if torch_dtype is not None:
+        kwargs["torch_dtype"] = torch_dtype
+
+    if model_name == "sd1.5":
+        return diffusers.models.AutoencoderKL.from_pretrained(base_path, **kwargs)
+
+    if model_name in {"flux", "flux1"}:
+        return diffusers.models.AutoencoderKL.from_pretrained(
+            base_path,
+            subfolder="vae",
+            **kwargs,
+        )
+
+    if model_name == "flux2":
+        return diffusers.models.AutoencoderKLFlux2.from_pretrained(
+            base_path,
+            subfolder="vae",
+            **kwargs,
+        )
+
+    raise ValueError(
+        f"Unsupported texture VAE model_name={model_name}. "
+        "Supported values: sd1.5, flux1, flux2."
+    )

sefi/pipeline.py

@@ -0,0 +1,161 @@
+"""SeFi-Image inference pipeline wrapper."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Iterable
+
+import torch
+from PIL import Image
+
+from .checkpoints import (
+    resolve_config_path,
+    resolve_checkpoint_to_local,
+)
+from .registry import ModelSpec, infer_model_spec
+from .resolution import resolve_image_size
+from .runtime import load_runtime_symbols
+
+
+SUPPORTED_DISTILL_STEPS = {4, 8, 10}
+
+
+class SEFIInferencePipeline:
+    """Inference wrapper for SeFi-Image checkpoints."""
+
+    def __init__(
+        self,
+        *,
+        spec: ModelSpec,
+        runner,
+        checkpoint_path: str,
+        checkpoint_uri: str,
+    ) -> None:
+        self.spec = spec
+        self.runner = runner
+        self.checkpoint_path = checkpoint_path
+        self.checkpoint_uri = checkpoint_uri
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        checkpoint: str,
+        *,
+        cache_dir: str | Path = "outputs/model_weights/sefi_inference",
+        config: str | Path | None = None,
+        device: str | None = None,
+        dtype: str | None = None,
+        delta_t: float | None = None,
+        timestep_shift_alpha: float | None = None,
+        debug_assert_schedule: bool = False,
+        autoguidance_config: str | None = None,
+        autoguidance_checkpoint: str | None = None,
+        guidance_interval_sigma_lo: float | None = None,
+        guidance_interval_sigma_hi: float | None = None,
+    ) -> "SEFIInferencePipeline":
+        runner_cls, load_config = load_runtime_symbols()
+
+        local_checkpoint, checkpoint_uri = resolve_checkpoint_to_local(
+            checkpoint=checkpoint,
+            cache_dir=cache_dir,
+        )
+        resolved_config = load_config(resolve_config_path(local_checkpoint, config))
+        spec = infer_model_spec(
+            resolved_config,
+            checkpoint_uri=checkpoint_uri,
+            checkpoint_path=local_checkpoint,
+        )
+        resolved_device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        resolved_dtype = dtype or spec.default_dtype
+        resolved_delta_t = delta_t if delta_t is not None else spec.default_delta_t
+        resolved_timestep_shift_alpha = (
+            timestep_shift_alpha
+            if timestep_shift_alpha is not None
+            else spec.default_timestep_shift_alpha
+        )
+
+        runner = runner_cls(
+            resolved_config,
+            checkpoint_path=local_checkpoint,
+            device=resolved_device,
+            debug_assert_schedule=debug_assert_schedule,
+            delta_t_override=resolved_delta_t,
+            inference_dtype=resolved_dtype,
+            timestep_shift_alpha=resolved_timestep_shift_alpha,
+            autoguidance_config_path=autoguidance_config,
+            autoguidance_checkpoint_path=autoguidance_checkpoint,
+            guidance_interval_sigma_lo=guidance_interval_sigma_lo,
+            guidance_interval_sigma_hi=guidance_interval_sigma_hi,
+        )
+
+        return cls(
+            spec=spec,
+            runner=runner,
+            checkpoint_path=local_checkpoint,
+            checkpoint_uri=checkpoint_uri,
+        )
+
+    def __call__(
+        self,
+        prompts: str | Iterable[str],
+        *,
+        num_inference_steps: int | None = None,
+        guidance_scale: float | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        batch_size: int | None = None,
+        seed: int | None = None,
+        generator: torch.Generator | None = None,
+    ) -> list[Image.Image]:
+        prompt_list = [prompts] if isinstance(prompts, str) else list(prompts)
+        if not prompt_list:
+            return []
+
+        steps = int(
+            num_inference_steps
+            if num_inference_steps is not None
+            else self.spec.default_steps
+        )
+        guidance = float(
+            guidance_scale
+            if guidance_scale is not None
+            else self.spec.default_guidance_scale
+        )
+        size = resolve_image_size(
+            height=height,
+            width=width,
+            default_height=self.spec.default_height,
+            default_width=self.spec.default_width,
+        )
+
+        if self.spec.is_distilled:
+            if steps not in SUPPORTED_DISTILL_STEPS:
+                raise ValueError(
+                    "SEFI Turbo models currently support "
+                    f"{sorted(SUPPORTED_DISTILL_STEPS)} steps, got {steps}."
+                )
+            if guidance != 1.0:
+                raise ValueError("SEFI Turbo models should run with guidance_scale=1.0.")
+
+        bs = int(batch_size or len(prompt_list))
+        if bs <= 0:
+            raise ValueError("batch_size must be > 0.")
+
+        gen = generator
+        if gen is None and seed is not None:
+            gen = torch.Generator(device=str(self.runner.device)).manual_seed(int(seed))
+
+        images: list[Image.Image] = []
+        for start in range(0, len(prompt_list), bs):
+            chunk = prompt_list[start : start + bs]
+            images.extend(
+                self.runner.generate_batch(
+                    prompts=chunk,
+                    num_inference_steps=steps,
+                    guidance_scale=guidance,
+                    height=size.height,
+                    width=size.width,
+                    generator=gen,
+                )
+            )
+        return images

sefi/registry.py

@@ -0,0 +1,203 @@
+"""Checkpoint-derived model metadata for SeFi-Image inference."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+from omegaconf import OmegaConf
+
+
+ModelFamily = Literal["base", "rl", "turbo"]
+
+
+@dataclass(frozen=True)
+class ModelSpec:
+    name: str
+    family: ModelFamily
+    scale: str
+    default_height: int = 1024
+    default_width: int = 1024
+    default_steps: int = 50
+    default_guidance_scale: float = 4.0
+    default_delta_t: float | None = None
+    default_timestep_shift_alpha: float = 1.0
+    default_dtype: str = "bf16"
+
+    @property
+    def is_distilled(self) -> bool:
+        return self.family == "turbo"
+
+
+def _string_option(config, *keys: str) -> str:
+    for key in keys:
+        value = OmegaConf.select(config, key, default=None)
+        if value is not None:
+            text = str(value).strip()
+            if text:
+                return text
+    return ""
+
+
+def _int_option(config, *keys: str, default: int) -> int:
+    for key in keys:
+        value = OmegaConf.select(config, key, default=None)
+        if value is not None:
+            return int(value)
+    return int(default)
+
+
+def _float_option(config, *keys: str, default: float | None) -> float | None:
+    for key in keys:
+        value = OmegaConf.select(config, key, default=None)
+        if value is not None:
+            return float(value)
+    return default
+
+
+def _checkpoint_hint(checkpoint_uri: str, checkpoint_path: str) -> str:
+    parts = [checkpoint_uri, checkpoint_path]
+    path = Path(checkpoint_path)
+    parts.extend(str(part) for part in path.parts[-4:])
+    return " ".join(parts).lower()
+
+
+def _normalize_family(value: str) -> ModelFamily | None:
+    text = value.strip().lower().replace("_", "-")
+    if text in {"base", "sft"}:
+        return "base"
+    if text in {"rl", "reward", "posttrain", "post-training"}:
+        return "rl"
+    if text in {"turbo", "distill", "distilled", "dmd", "dmd2"}:
+        return "turbo"
+    return None
+
+
+def _infer_family(config, checkpoint_uri: str, checkpoint_path: str) -> ModelFamily:
+    configured = _string_option(
+        config,
+        "inference.family",
+        "inference.variant",
+        "model.family",
+        "model.variant",
+    )
+    if configured:
+        family = _normalize_family(configured)
+        if family is None:
+            raise ValueError(
+                "Unsupported SeFi-Image model family in config: "
+                f"{configured}. Expected base, rl, or turbo."
+            )
+        return family
+
+    hint = _checkpoint_hint(checkpoint_uri, checkpoint_path)
+    normalized = re.sub(r"[^a-z0-9]+", "-", hint)
+    if "turbo" in normalized or "distill" in normalized or "dmd" in normalized:
+        return "turbo"
+    if re.search(r"(^|-)rl($|-)", normalized) or "scalar" in normalized:
+        return "rl"
+    if "base" in normalized or "sft" in normalized:
+        return "base"
+
+    raise ValueError(
+        "Could not infer SeFi-Image checkpoint family from checkpoint name. "
+        "Use a checkpoint path or Hugging Face repo id containing Base, RL, or "
+        "Turbo, or add inference.family to sefi_config.yaml."
+    )
+
+
+def _infer_scale(config, checkpoint_uri: str, checkpoint_path: str) -> str:
+    configured = _string_option(config, "model.transformer_scale")
+    if configured and configured != "custom":
+        return configured.lower()
+
+    model_name = _string_option(config, "model.model_name").lower()
+    match = re.search(r"([0-9]+(?:p[0-9]+)?b)", model_name)
+    if match:
+        return match.group(1)
+
+    hint = _checkpoint_hint(checkpoint_uri, checkpoint_path)
+    match = re.search(r"([0-9]+(?:p[0-9]+)?b)", hint)
+    if match:
+        return match.group(1).lower()
+
+    raise ValueError(
+        "Could not infer SeFi-Image model scale from config or checkpoint name."
+    )
+
+
+def _default_name(family: ModelFamily, scale: str) -> str:
+    public_scale = scale.upper().replace("P", ".")
+    suffix = {"base": "Base", "rl": "RL", "turbo": "turbo"}[family]
+    return f"SeFi-Image-{public_scale}-{suffix}"
+
+
+def _default_steps(family: ModelFamily) -> int:
+    return 4 if family == "turbo" else 50
+
+
+def _default_guidance_scale(family: ModelFamily) -> float:
+    return 1.0 if family == "turbo" else 4.0
+
+
+def _default_timestep_shift_alpha(family: ModelFamily) -> float:
+    return 0.3 if family in {"base", "rl"} else 1.0
+
+
+def _default_dtype(config) -> str:
+    dtype = _string_option(config, "inference.dtype", "training.mixed_precision")
+    dtype = dtype.lower()
+    if dtype in {"bf16", "bfloat16"}:
+        return "bf16"
+    if dtype in {"fp32", "float32", "no", "none"}:
+        return "fp32"
+    return "bf16"
+
+
+def infer_model_spec(
+    config,
+    *,
+    checkpoint_uri: str,
+    checkpoint_path: str,
+) -> ModelSpec:
+    family = _infer_family(config, checkpoint_uri, checkpoint_path)
+    scale = _infer_scale(config, checkpoint_uri, checkpoint_path)
+    resolution = _int_option(config, "data.resolution", default=1024)
+    height = _int_option(config, "inference.height", "data.height", default=resolution)
+    width = _int_option(config, "inference.width", "data.width", default=resolution)
+    name = _string_option(config, "inference.model_name", "model.display_name")
+
+    return ModelSpec(
+        name=name or _default_name(family, scale),
+        family=family,
+        scale=scale,
+        default_height=height,
+        default_width=width,
+        default_steps=_int_option(
+            config,
+            "inference.steps",
+            "inference.default_steps",
+            default=_default_steps(family),
+        ),
+        default_guidance_scale=_float_option(
+            config,
+            "inference.guidance_scale",
+            "inference.default_guidance_scale",
+            default=_default_guidance_scale(family),
+        )
+        or _default_guidance_scale(family),
+        default_delta_t=_float_option(
+            config,
+            "inference.delta_t",
+            default=None,
+        ),
+        default_timestep_shift_alpha=_float_option(
+            config,
+            "inference.timestep_shift_alpha",
+            default=_default_timestep_shift_alpha(family),
+        )
+        or _default_timestep_shift_alpha(family),
+        default_dtype=_default_dtype(config),
+    )

sefi/resolution.py

@@ -0,0 +1,32 @@
+"""Resolution helpers for SEFI inference."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class ImageSize:
+    height: int
+    width: int
+
+
+def resolve_image_size(
+    *,
+    height: int | None,
+    width: int | None,
+    default_height: int,
+    default_width: int,
+) -> ImageSize:
+    resolved = ImageSize(
+        height=int(height if height is not None else default_height),
+        width=int(width if width is not None else default_width),
+    )
+    if resolved.height <= 0 or resolved.width <= 0:
+        raise ValueError(f"Image size must be positive, got {resolved}.")
+    if resolved.height % 16 != 0 or resolved.width % 16 != 0:
+        raise ValueError(
+            "SEFI image size must be divisible by 16, "
+            f"got height={resolved.height}, width={resolved.width}."
+        )
+    return resolved

sefi/runner.py

@@ -0,0 +1,793 @@
+"""SEFI T2I inference runner with three-phase masked denoising."""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+from typing import Optional
+
+import torch
+from PIL import Image
+from torch import Tensor
+
+from .builder import (
+    build_components,
+    build_lightweight_transformer,
+    _derive_semantic_channels,
+    _derive_text_output_dim,
+    _derive_texture_channels,
+    text_encoder_signature,
+)
+from .config import load_config
+from .modeling import Qwen3VLTextEncoder
+
+
+def _resolve_weight_dtype(config, *, override: Optional[str] = None) -> torch.dtype:
+    if override is not None:
+        normalized = str(override).strip().lower()
+        if normalized == "bf16":
+            return torch.bfloat16
+        if normalized in {"fp32", "float32"}:
+            return torch.float32
+        raise ValueError(
+            f"Unsupported inference dtype: {override}. Expected one of ['bf16', 'fp32']."
+        )
+
+    precision = str(getattr(config.training, "mixed_precision", "bf16")).lower()
+    if precision == "fp16":
+        return torch.float16
+    if precision in {"fp32", "float32", "no"}:
+        return torch.float32
+    return torch.bfloat16
+
+
+def _training_sefi_cfg(config):
+    cfg = config.training.get("sefi", None)
+    if cfg is not None:
+        return cfg
+    raise ValueError("Config requires training.sefi section.")
+
+
+def _apply_timestep_shift_unit_interval(u_unit: Tensor, alpha: float) -> Tensor:
+    """Apply t' = alpha*t / (1 + (alpha-1)*t) on unit coordinate u in [0, 1]."""
+    alpha = float(alpha)
+    if alpha <= 0:
+        raise ValueError(f"timestep_shift_alpha must be > 0, got {alpha}")
+    if alpha == 1.0:
+        return u_unit
+    denominator = 1.0 + (alpha - 1.0) * u_unit
+    return (alpha * u_unit) / denominator
+
+
+def _combine_guided_velocity(base_pred: Tensor, cond_pred: Tensor, guidance_scale: float) -> Tensor:
+    """Shared guidance formula: base + scale * (conditioned - base)."""
+    return base_pred + float(guidance_scale) * (cond_pred - base_pred)
+
+
+def _resolve_guidance_interval_sigma(
+    sigma_lo: Optional[float],
+    sigma_hi: Optional[float],
+) -> tuple[Optional[float], Optional[float]]:
+    if sigma_lo is None and sigma_hi is None:
+        return None, None
+    if sigma_lo is None or sigma_hi is None:
+        raise ValueError(
+            "Limited interval guidance requires both "
+            "guidance_interval_sigma_lo and guidance_interval_sigma_hi, or neither."
+        )
+
+    sigma_lo = float(sigma_lo)
+    sigma_hi = float(sigma_hi)
+    if not math.isfinite(sigma_lo) or not math.isfinite(sigma_hi):
+        raise ValueError("guidance interval sigma thresholds must be finite.")
+    if sigma_lo < 0.0 or sigma_hi < 0.0:
+        raise ValueError("guidance interval sigma thresholds must be >= 0.")
+    if sigma_lo >= sigma_hi:
+        raise ValueError("guidance_interval_sigma_lo must be < guidance_interval_sigma_hi.")
+    return sigma_lo, sigma_hi
+
+
+def _guidance_interval_is_active(
+    sigma: Tensor | float,
+    sigma_lo: Optional[float],
+    sigma_hi: Optional[float],
+) -> bool:
+    if sigma_lo is None and sigma_hi is None:
+        return True
+    if sigma_lo is None or sigma_hi is None:
+        raise ValueError("guidance interval sigma bounds must be paired.")
+
+    sigma_value = float(sigma.item()) if isinstance(sigma, Tensor) else float(sigma)
+    return float(sigma_lo) < sigma_value <= float(sigma_hi)
+
+
+def _normalize_optional_path(path: Optional[str]) -> str:
+    if path is None:
+        return ""
+    return str(path).strip()
+
+
+def _resolve_autoguidance_paths(
+    autoguidance_config_path: Optional[str],
+    autoguidance_checkpoint_path: Optional[str],
+) -> tuple[str, str]:
+    config_path = _normalize_optional_path(autoguidance_config_path)
+    checkpoint_path = _normalize_optional_path(autoguidance_checkpoint_path)
+    if bool(config_path) != bool(checkpoint_path):
+        raise ValueError(
+            "AutoGuidance requires both --autoguidance_config and "
+            "--autoguidance_checkpoint, or neither."
+        )
+    return config_path, checkpoint_path
+
+
+def _validate_autoguidance_guidance_scale(enabled: bool, guidance_scale: float) -> None:
+    if enabled and float(guidance_scale) <= 1.0:
+        raise ValueError("AutoGuidance requires guidance_scale > 1.0.")
+
+
+def _resolve_checkpoint_file(checkpoint_path: str) -> str:
+    if os.path.isdir(checkpoint_path):
+        transformer_dir = os.path.join(checkpoint_path, "transformer")
+        sharded_safetensors = os.path.join(
+            transformer_dir,
+            "diffusion_pytorch_model.safetensors.index.json",
+        )
+        if os.path.isfile(sharded_safetensors):
+            return sharded_safetensors
+
+        safetensors_state = os.path.join(
+            transformer_dir,
+            "diffusion_pytorch_model.safetensors",
+        )
+        if os.path.isfile(safetensors_state):
+            return safetensors_state
+
+        torch_state = os.path.join(transformer_dir, "diffusion_pytorch_model.bin")
+        if os.path.isfile(torch_state):
+            return torch_state
+
+        raise FileNotFoundError(
+            f"Unsupported SEFI inference checkpoint directory: {checkpoint_path}. "
+            "Expected transformer/diffusion_pytorch_model.safetensors or "
+            "transformer/diffusion_pytorch_model.safetensors.index.json."
+        )
+
+    if not os.path.exists(checkpoint_path):
+        raise FileNotFoundError(f"Checkpoint path not found: {checkpoint_path}")
+
+    return checkpoint_path
+
+
+def _extract_state_dict(checkpoint: dict) -> dict:
+    if not isinstance(checkpoint, dict):
+        raise ValueError("Checkpoint must be a dict-like object.")
+
+    if "model_state_dict" in checkpoint and isinstance(checkpoint["model_state_dict"], dict):
+        return checkpoint["model_state_dict"]
+    if "module" in checkpoint and isinstance(checkpoint["module"], dict):
+        return checkpoint["module"]
+    if "state_dict" in checkpoint and isinstance(checkpoint["state_dict"], dict):
+        return checkpoint["state_dict"]
+    if checkpoint and all(isinstance(v, torch.Tensor) for v in checkpoint.values()):
+        return checkpoint
+
+    raise ValueError(
+        "Unsupported checkpoint format. Expected one of: "
+        "model_state_dict / module / state_dict / plain state_dict."
+    )
+
+
+def _load_checkpoint_payload(checkpoint_file: str):
+    if checkpoint_file.endswith(".safetensors.index.json"):
+        from safetensors.torch import load_file
+
+        with open(checkpoint_file, "r", encoding="utf-8") as handle:
+            index = json.load(handle)
+        weight_map = index.get("weight_map", None)
+        if not isinstance(weight_map, dict) or not weight_map:
+            raise ValueError(f"Invalid safetensors index file: {checkpoint_file}")
+
+        base_dir = os.path.dirname(checkpoint_file)
+        state_dict = {}
+        for shard_name in sorted(set(weight_map.values())):
+            shard_path = os.path.join(base_dir, shard_name)
+            if not os.path.isfile(shard_path):
+                raise FileNotFoundError(f"Missing safetensors shard: {shard_path}")
+            state_dict.update(load_file(shard_path))
+        return state_dict
+
+    if checkpoint_file.endswith(".safetensors"):
+        from safetensors.torch import load_file
+
+        return load_file(checkpoint_file)
+
+    return torch.load(checkpoint_file, map_location="cpu")
+
+
+def _strip_prefix_if_needed(state_dict: dict, prefix: str) -> dict:
+    if state_dict and all(k.startswith(prefix) for k in state_dict):
+        return {k[len(prefix) :]: v for k, v in state_dict.items()}
+    return state_dict
+
+
+def _load_transformer_state_dict_strict_shapes(
+    transformer,
+    checkpoint_path: str,
+    *,
+    label: str,
+) -> None:
+    checkpoint_file = _resolve_checkpoint_file(checkpoint_path)
+    print(f"Loading {label} checkpoint from {checkpoint_file}")
+    payload = _load_checkpoint_payload(checkpoint_file)
+    state_dict = _extract_state_dict(payload)
+    state_dict = _strip_prefix_if_needed(state_dict, "module.")
+
+    target_state = transformer.state_dict()
+    compatible_state = {}
+    mismatched = []
+    for key, value in state_dict.items():
+        if key not in target_state:
+            continue
+        if tuple(value.shape) != tuple(target_state[key].shape):
+            mismatched.append(
+                f"{key}: checkpoint={tuple(value.shape)} vs model={tuple(target_state[key].shape)}"
+            )
+            continue
+        compatible_state[key] = value
+
+    if mismatched:
+        raise ValueError(f"{label} checkpoint has shape-mismatched keys: {mismatched[:10]}")
+    if not compatible_state:
+        raise ValueError(
+            f"{label} checkpoint has zero loadable parameters for the constructed model: "
+            f"{checkpoint_path}"
+        )
+
+    missing, unexpected = transformer.load_state_dict(compatible_state, strict=False)
+    if missing:
+        print(f"  Warning - {label} missing keys: {missing[:10]}")
+    if unexpected:
+        print(f"  Warning - {label} unexpected keys: {unexpected[:10]}")
+
+
+class SEFIInferenceRunner:
+    """Inference runner for SEFI-T2I with three-phase masked denoising."""
+
+    def __init__(
+        self,
+        config,
+        *,
+        checkpoint_path: str = "",
+        device: str = "cuda",
+        debug_assert_schedule: bool = False,
+        delta_t_override: Optional[float] = None,
+        inference_dtype: Optional[str] = None,
+        timestep_shift_alpha: float = 1.0,
+        autoguidance_config_path: Optional[str] = None,
+        autoguidance_checkpoint_path: Optional[str] = None,
+        guidance_interval_sigma_lo: Optional[float] = None,
+        guidance_interval_sigma_hi: Optional[float] = None,
+    ):
+        from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor
+
+        self.config = config
+        self.device = torch.device(device)
+        self.component_dtype = _resolve_weight_dtype(config)
+        self.weight_dtype = _resolve_weight_dtype(config, override=inference_dtype)
+        (
+            self.autoguidance_config_path,
+            self.autoguidance_checkpoint_path,
+        ) = _resolve_autoguidance_paths(
+            autoguidance_config_path,
+            autoguidance_checkpoint_path,
+        )
+        self.autoguidance_enabled = bool(self.autoguidance_config_path)
+        self.autoguidance_transformer = None
+        self.autoguidance_text_encoder = None
+        self.autoguidance_reuse_main_text_encoder = True
+        (
+            self.guidance_interval_sigma_lo,
+            self.guidance_interval_sigma_hi,
+        ) = _resolve_guidance_interval_sigma(
+            guidance_interval_sigma_lo,
+            guidance_interval_sigma_hi,
+        )
+        self.guidance_interval_enabled = self.guidance_interval_sigma_lo is not None
+
+        components = build_components(config, component_dtype=self.component_dtype)
+        self.transformer = components.transformer.to(
+            device=self.device,
+            dtype=self.weight_dtype,
+        ).eval()
+        self.text_encoder = components.text_encoder.to(
+            device=self.device,
+            dtype=self.component_dtype,
+        ).eval()
+        self.texture_codec = components.texture_codec.to(
+            device=self.device,
+            dtype=self.component_dtype,
+        ).eval()
+        self.noise_scheduler = components.noise_scheduler
+        self.pipeline_cls = components.pipeline_cls
+        self.semantic_channels = int(components.semantic_channels)
+        self.texture_channels = int(components.texture_channels)
+        self.total_channels = int(components.total_channels)
+
+        self.debug_assert_schedule = bool(debug_assert_schedule)
+        self.timestep_shift_alpha = float(timestep_shift_alpha)
+        if self.timestep_shift_alpha <= 0:
+            raise ValueError(
+                "timestep_shift_alpha must be > 0. "
+                f"Got {self.timestep_shift_alpha}."
+            )
+        self._configure_delta_t(delta_t_override)
+        shift_enabled = self.timestep_shift_alpha != 1.0
+        print(
+            "Inference timestep schedule: "
+            f"timestep_shift_alpha={self.timestep_shift_alpha:.6f}, "
+            f"delta_t={self.delta_t:.6f}, shift_enabled={shift_enabled}"
+        )
+        if self.guidance_interval_enabled:
+            print(
+                "Limited interval guidance enabled on base sigma: "
+                f"({self.guidance_interval_sigma_lo:.6f}, "
+                f"{self.guidance_interval_sigma_hi:.6f}]"
+            )
+
+        texture_vae_cfg = self.texture_codec.texture_vae.config
+        self.vae_scale_factor = 2 ** (len(texture_vae_cfg.block_out_channels) - 1)
+        self.image_processor = Flux2ImageProcessor(
+            vae_scale_factor=self.vae_scale_factor * 2
+        )
+
+        for module in (self.transformer, self.text_encoder, self.texture_codec):
+            for param in module.parameters():
+                param.requires_grad = False
+
+        if checkpoint_path:
+            self.load_checkpoint(checkpoint_path)
+
+        if self.autoguidance_enabled:
+            self._load_autoguidance_model()
+
+    def _configure_delta_t(self, delta_t_override: Optional[float]) -> None:
+        sefi_cfg = _training_sefi_cfg(self.config)
+
+        delta_t_min_raw = sefi_cfg.get("delta_t_min", None)
+        delta_t_max_raw = sefi_cfg.get("delta_t_max", None)
+        if delta_t_min_raw is None or delta_t_max_raw is None:
+            raise ValueError("training.sefi.delta_t_min and delta_t_max are required.")
+
+        self.delta_t_min = float(delta_t_min_raw)
+        self.delta_t_max = float(delta_t_max_raw)
+        if self.delta_t_min < 0 or self.delta_t_min > 1:
+            raise ValueError("training.sefi.delta_t_min must be in [0, 1].")
+        if self.delta_t_max < 0 or self.delta_t_max > 1:
+            raise ValueError("training.sefi.delta_t_max must be in [0, 1].")
+        if self.delta_t_min > self.delta_t_max:
+            raise ValueError("training.sefi.delta_t_min must be <= delta_t_max.")
+
+        if delta_t_override is None:
+            self.delta_t = self.delta_t_max
+            print(
+                "Warning: --delta-t not provided. "
+                f"Using training.sefi.delta_t_max={self.delta_t_max:.6f} for inference."
+            )
+            return
+
+        self.delta_t = float(delta_t_override)
+        if self.delta_t < 0 or self.delta_t > 1:
+            raise ValueError("inference delta_t must be in [0, 1].")
+        if self.delta_t < self.delta_t_min or self.delta_t > self.delta_t_max:
+            print(
+                "Warning: inference delta_t is outside training range "
+                f"[{self.delta_t_min:.6f}, {self.delta_t_max:.6f}]. "
+                f"Got delta_t={self.delta_t:.6f}."
+            )
+
+    def load_checkpoint(self, checkpoint_path: str):
+        ckpt_file = _resolve_checkpoint_file(checkpoint_path)
+        print(f"Loading checkpoint from {ckpt_file}")
+        ckpt = _load_checkpoint_payload(ckpt_file)
+        state_dict = _extract_state_dict(ckpt)
+        state_dict = _strip_prefix_if_needed(state_dict, "module.")
+
+        missing, unexpected = self.transformer.load_state_dict(state_dict, strict=False)
+        if missing:
+            raise ValueError(f"Checkpoint is missing transformer keys: {missing[:10]}")
+        if unexpected:
+            raise ValueError(f"Checkpoint has unexpected transformer keys: {unexpected[:10]}")
+
+    def _load_autoguidance_model(self) -> None:
+        autoguidance_config = load_config(self.autoguidance_config_path)
+        self.autoguidance_config = autoguidance_config
+
+        ag_semantic_channels = _derive_semantic_channels(autoguidance_config)
+        ag_texture_channels = _derive_texture_channels(autoguidance_config)
+        if ag_semantic_channels != self.semantic_channels:
+            raise ValueError(
+                "AutoGuidance semantic channel mismatch: "
+                f"main={self.semantic_channels}, small={ag_semantic_channels}."
+            )
+        if ag_texture_channels != self.texture_channels:
+            raise ValueError(
+                "AutoGuidance texture channel mismatch: "
+                f"main={self.texture_channels}, small={ag_texture_channels}."
+            )
+
+        ag_text_output_dim = _derive_text_output_dim(autoguidance_config)
+        autoguidance_transformer = build_lightweight_transformer(
+            autoguidance_config,
+            total_channels=self.total_channels,
+            text_output_dim=ag_text_output_dim,
+        )
+        _load_transformer_state_dict_strict_shapes(
+            autoguidance_transformer,
+            self.autoguidance_checkpoint_path,
+            label="AutoGuidance",
+        )
+        self.autoguidance_transformer = autoguidance_transformer.to(
+            device=self.device,
+            dtype=self.weight_dtype,
+        ).eval()
+        for param in self.autoguidance_transformer.parameters():
+            param.requires_grad = False
+
+        self.autoguidance_reuse_main_text_encoder = (
+            text_encoder_signature(self.config) == text_encoder_signature(autoguidance_config)
+        )
+        if self.autoguidance_reuse_main_text_encoder:
+            print("AutoGuidance reuses main prompt embeddings.")
+        else:
+            text_cfg = autoguidance_config.model.text_encoder
+            self.autoguidance_text_encoder = Qwen3VLTextEncoder(
+                model_name=str(text_cfg.model_name),
+                weights_root=str(text_cfg.get("weights_root", "outputs/model_weights")),
+                max_length=int(text_cfg.max_length),
+                hidden_layers=[int(x) for x in text_cfg.hidden_layers],
+                torch_dtype=self.component_dtype,
+            ).to(device=self.device, dtype=self.component_dtype).eval()
+            if int(self.autoguidance_text_encoder.output_dim) != int(ag_text_output_dim):
+                raise ValueError(
+                    "AutoGuidance text encoder output dim mismatch: "
+                    f"loaded={self.autoguidance_text_encoder.output_dim}, "
+                    f"expected={ag_text_output_dim}."
+                )
+            for param in self.autoguidance_text_encoder.parameters():
+                param.requires_grad = False
+            print("AutoGuidance uses a separate small-model text encoder.")
+
+        print(
+            "Loaded AutoGuidance model: "
+            f"config={self.autoguidance_config_path}, "
+            f"checkpoint={self.autoguidance_checkpoint_path}"
+        )
+
+    def _timesteps_and_sigmas(
+        self,
+        u_continuous: Tensor,
+        *,
+        n_dim: int,
+        dtype: torch.dtype,
+    ) -> tuple[Tensor, Tensor]:
+        num_steps = int(self.noise_scheduler.config.num_train_timesteps)
+        indices = (u_continuous * (num_steps - 1)).long().clamp(0, num_steps - 1)
+
+        timesteps = self.noise_scheduler.timesteps[indices.cpu()].to(self.device)
+        sigmas = self.noise_scheduler.sigmas[indices.cpu()].to(
+            device=self.device,
+            dtype=dtype,
+        )
+        while sigmas.ndim < n_dim:
+            sigmas = sigmas.unsqueeze(-1)
+        return timesteps, sigmas
+
+    def _assert_shifted_schedule(
+        self,
+        u_base_unit: Tensor,
+        u_sem_raw_schedule: Tensor,
+        eps: float = 1e-6,
+    ) -> None:
+        if u_base_unit.ndim != 1 or u_sem_raw_schedule.ndim != 1:
+            raise ValueError("u_base_unit and u_sem_raw_schedule must be 1D tensors.")
+        if u_base_unit.shape != u_sem_raw_schedule.shape:
+            raise ValueError("u_base_unit and u_sem_raw_schedule must have the same shape.")
+
+        expected_u_max = 1.0 + self.delta_t
+        if abs(float(u_base_unit[0].item()) - 0.0) > eps:
+            raise ValueError(
+                f"Invalid u_base_unit[0], expected 0, got {float(u_base_unit[0].item()):.6f}"
+            )
+        if abs(float(u_base_unit[-1].item()) - 1.0) > eps:
+            raise ValueError(
+                f"Invalid u_base_unit[-1], expected 1, got {float(u_base_unit[-1].item()):.6f}"
+            )
+        if abs(float(u_sem_raw_schedule[0].item()) - 0.0) > eps:
+            raise ValueError(
+                "Invalid shifted schedule start, expected 0, "
+                f"got {float(u_sem_raw_schedule[0].item()):.6f}"
+            )
+        if abs(float(u_sem_raw_schedule[-1].item()) - expected_u_max) > eps:
+            raise ValueError(
+                "Invalid shifted schedule end, expected 1+delta_t, "
+                f"got {float(u_sem_raw_schedule[-1].item()):.6f}, "
+                f"expected={expected_u_max:.6f}"
+            )
+
+        diffs = u_sem_raw_schedule[1:] - u_sem_raw_schedule[:-1]
+        if torch.any(diffs < -eps):
+            index = int(torch.nonzero(diffs < -eps, as_tuple=False)[0, 0].item())
+            raise ValueError(
+                "Shifted u_sem_raw schedule must be monotonic non-decreasing, "
+                f"but got decrease at step {index}: "
+                f"{float(u_sem_raw_schedule[index].item()):.6f} -> "
+                f"{float(u_sem_raw_schedule[index + 1].item()):.6f}"
+            )
+
+    def _assert_dual_time_invariants(
+        self,
+        u_sem: Tensor,
+        u_tex: Tensor,
+        sigmas_sem: Tensor,
+        sigmas_tex: Tensor,
+        eps: float = 1e-6,
+    ) -> None:
+        u_violation = u_sem < u_tex
+        if torch.any(u_violation):
+            index = int(torch.nonzero(u_violation, as_tuple=False)[0, 0].item())
+            raise ValueError(
+                "Dual-time invariant violated: expected u_sem >= u_tex, got "
+                f"u_sem[{index}]={float(u_sem[index].item()):.6f}, "
+                f"u_tex[{index}]={float(u_tex[index].item()):.6f}."
+            )
+
+        sigma_violation = sigmas_sem > (sigmas_tex + eps)
+        if torch.any(sigma_violation):
+            index = int(torch.nonzero(sigma_violation, as_tuple=False)[0, 0].item())
+            sigma_sem_flat = sigmas_sem.reshape(sigmas_sem.shape[0], -1)
+            sigma_tex_flat = sigmas_tex.reshape(sigmas_tex.shape[0], -1)
+            raise ValueError(
+                "Dual-time invariant violated: expected sigmas_sem <= sigmas_tex, got "
+                f"sigmas_sem[{index}]={float(sigma_sem_flat[index, 0].item()):.6f}, "
+                f"sigmas_tex[{index}]={float(sigma_tex_flat[index, 0].item()):.6f}."
+            )
+
+    def _prepare_latents(
+        self,
+        *,
+        batch_size: int,
+        height: int,
+        width: int,
+        generator: Optional[torch.Generator],
+    ) -> tuple[Tensor, Tensor, int, int]:
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+        latents = torch.randn(
+            (batch_size, self.total_channels, height // 2, width // 2),
+            generator=generator,
+            device=self.device,
+            dtype=self.weight_dtype,
+        )
+        latent_ids = self.pipeline_cls._prepare_latent_ids(latents).to(self.device)
+        return latents, latent_ids, height, width
+
+    def _predict_velocity(
+        self,
+        transformer,
+        *,
+        packed_latents: Tensor,
+        timesteps_sem: Tensor,
+        timesteps_tex: Tensor,
+        encoder_hidden_states: Tensor,
+        txt_ids: Tensor,
+        img_ids: Tensor,
+    ) -> Tensor:
+        pred = transformer(
+            hidden_states=packed_latents,
+            timestep_sem=timesteps_sem / 1000,
+            timestep_tex=timesteps_tex / 1000,
+            encoder_hidden_states=encoder_hidden_states,
+            txt_ids=txt_ids,
+            img_ids=img_ids,
+        )
+        pred = pred[:, : packed_latents.size(1)]
+        return self.pipeline_cls._unpack_latents_with_ids(pred, img_ids)
+
+    @torch.no_grad()
+    def generate_batch(
+        self,
+        *,
+        prompts: list[str],
+        num_inference_steps: int,
+        guidance_scale: float,
+        height: int,
+        width: int,
+        generator: Optional[torch.Generator] = None,
+    ) -> list[Image.Image]:
+        if num_inference_steps <= 0:
+            raise ValueError("num_inference_steps must be > 0")
+
+        batch_size = len(prompts)
+        if batch_size == 0:
+            return []
+
+        prompt_embeds, text_ids = self.text_encoder.encode(prompts, dtype=self.weight_dtype)
+        _validate_autoguidance_guidance_scale(
+            self.autoguidance_enabled,
+            guidance_scale,
+        )
+
+        if self.autoguidance_enabled:
+            if self.autoguidance_reuse_main_text_encoder:
+                autoguidance_prompt_embeds = prompt_embeds
+                autoguidance_text_ids = text_ids
+            else:
+                autoguidance_prompt_embeds, autoguidance_text_ids = (
+                    self.autoguidance_text_encoder.encode(
+                        prompts,
+                        dtype=self.weight_dtype,
+                    )
+                )
+            neg_prompt_embeds = None
+            neg_text_ids = None
+        elif guidance_scale > 1.0:
+            neg_prompts = [""] * batch_size
+            neg_prompt_embeds, neg_text_ids = self.text_encoder.encode(
+                neg_prompts,
+                dtype=self.weight_dtype,
+            )
+            autoguidance_prompt_embeds = None
+            autoguidance_text_ids = None
+        else:
+            autoguidance_prompt_embeds = None
+            autoguidance_text_ids = None
+            neg_prompt_embeds = None
+            neg_text_ids = None
+
+        latents, latent_ids, _, _ = self._prepare_latents(
+            batch_size=batch_size,
+            height=height,
+            width=width,
+            generator=generator,
+        )
+
+        u_base_unit = torch.linspace(
+            0.0,
+            1.0,
+            steps=num_inference_steps + 1,
+            device=self.device,
+            dtype=torch.float32,
+        )
+        u_shifted_unit = _apply_timestep_shift_unit_interval(
+            u_base_unit,
+            self.timestep_shift_alpha,
+        )
+        _, base_sigmas_schedule = self._timesteps_and_sigmas(
+            u_shifted_unit,
+            n_dim=1,
+            dtype=torch.float32,
+        )
+        u_sem_raw_schedule = u_shifted_unit * (1.0 + self.delta_t)
+        if self.debug_assert_schedule:
+            self._assert_shifted_schedule(
+                u_base_unit=u_base_unit,
+                u_sem_raw_schedule=u_sem_raw_schedule,
+            )
+
+        for step in range(num_inference_steps):
+            u_sem_raw_cur = torch.full(
+                (batch_size,),
+                float(u_sem_raw_schedule[step].item()),
+                device=self.device,
+            )
+            u_sem_raw_next = torch.full(
+                (batch_size,),
+                float(u_sem_raw_schedule[step + 1].item()),
+                device=self.device,
+            )
+
+            u_tex_cur = torch.clamp(u_sem_raw_cur - self.delta_t, min=0.0, max=1.0)
+            u_sem_cur = torch.clamp(u_sem_raw_cur, max=1.0)
+            u_tex_next = torch.clamp(u_sem_raw_next - self.delta_t, min=0.0, max=1.0)
+            u_sem_next = torch.clamp(u_sem_raw_next, max=1.0)
+
+            timesteps_sem_cur, sigmas_sem_cur = self._timesteps_and_sigmas(
+                u_sem_cur,
+                n_dim=latents.ndim,
+                dtype=latents.dtype,
+            )
+            timesteps_tex_cur, sigmas_tex_cur = self._timesteps_and_sigmas(
+                u_tex_cur,
+                n_dim=latents.ndim,
+                dtype=latents.dtype,
+            )
+            _, sigmas_sem_next = self._timesteps_and_sigmas(
+                u_sem_next,
+                n_dim=latents.ndim,
+                dtype=latents.dtype,
+            )
+            _, sigmas_tex_next = self._timesteps_and_sigmas(
+                u_tex_next,
+                n_dim=latents.ndim,
+                dtype=latents.dtype,
+            )
+            if self.debug_assert_schedule:
+                self._assert_dual_time_invariants(
+                    u_sem_cur,
+                    u_tex_cur,
+                    sigmas_sem_cur,
+                    sigmas_tex_cur,
+                )
+
+            guidance_active = _guidance_interval_is_active(
+                base_sigmas_schedule[step],
+                self.guidance_interval_sigma_lo,
+                self.guidance_interval_sigma_hi,
+            )
+            packed_latents = self.pipeline_cls._pack_latents(latents)
+            pred_cond = self._predict_velocity(
+                self.transformer,
+                packed_latents=packed_latents,
+                timesteps_sem=timesteps_sem_cur,
+                timesteps_tex=timesteps_tex_cur,
+                encoder_hidden_states=prompt_embeds,
+                txt_ids=text_ids,
+                img_ids=latent_ids,
+            )
+
+            if not guidance_active:
+                velocity = pred_cond
+            elif self.autoguidance_enabled:
+                pred_base = self._predict_velocity(
+                    self.autoguidance_transformer,
+                    packed_latents=packed_latents,
+                    timesteps_sem=timesteps_sem_cur,
+                    timesteps_tex=timesteps_tex_cur,
+                    encoder_hidden_states=autoguidance_prompt_embeds,
+                    txt_ids=autoguidance_text_ids,
+                    img_ids=latent_ids,
+                )
+                velocity = _combine_guided_velocity(
+                    pred_base,
+                    pred_cond,
+                    guidance_scale,
+                )
+            elif guidance_scale > 1.0:
+                pred_uncond = self._predict_velocity(
+                    self.transformer,
+                    packed_latents=packed_latents,
+                    timesteps_sem=timesteps_sem_cur,
+                    timesteps_tex=timesteps_tex_cur,
+                    encoder_hidden_states=neg_prompt_embeds,
+                    txt_ids=neg_text_ids,
+                    img_ids=latent_ids,
+                )
+                velocity = _combine_guided_velocity(
+                    pred_uncond,
+                    pred_cond,
+                    guidance_scale,
+                )
+            else:
+                velocity = pred_cond
+
+            vel_sem = velocity[:, : self.semantic_channels]
+            vel_tex = velocity[:, self.semantic_channels :]
+            lat_sem = latents[:, : self.semantic_channels]
+            lat_tex = latents[:, self.semantic_channels :]
+
+            dt_sem = sigmas_sem_next - sigmas_sem_cur
+            dt_tex = sigmas_tex_next - sigmas_tex_cur
+
+            lat_sem = lat_sem + dt_sem * vel_sem
+            lat_tex = lat_tex + dt_tex * vel_tex
+            latents = torch.cat([lat_sem, lat_tex], dim=1)
+
+        texture_latents = latents[:, self.semantic_channels :]
+        decoded = self.texture_codec.decode_texture(
+            texture_latents.to(dtype=self.component_dtype),
+            pipeline_cls=self.pipeline_cls,
+        )
+        return self.image_processor.postprocess(decoded, output_type="pil")

sefi/runtime.py

@@ -0,0 +1,13 @@
+"""Runtime helpers for the bundled SEFI inference package."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .config import load_config
+
+
+def load_runtime_symbols() -> tuple[Any, Any]:
+    from .runner import SEFIInferenceRunner
+
+    return SEFIInferenceRunner, load_config