core_anim/wanvideo_loop_runtime.py

from __future__ import annotations

import base64
import copy
from datetime import datetime, timezone
import gc
import hashlib
import importlib
import json
import os
from pathlib import Path
import time
from typing import Callable

import numpy as np
import torch
from PIL import Image

try:
    from .wanvideo_wrapper_bridge import get_wrapper_node_class, load_wrapper_module
except ImportError:  # pragma: no cover
    from wanvideo_wrapper_bridge import get_wrapper_node_class, load_wrapper_module


def _assert_zerogpu_parent_fork_clean(event: str) -> None:
    if os.environ.get("SPACES_ZERO_GPU", "").strip().lower() not in {"1", "true", "yes", "on"}:
        return
    try:
        from spaces.zero import wrappers as zero_wrappers
        if zero_wrappers.forked:
            return
    except Exception:
        pass
    read_fd, write_fd = os.pipe()
    pid = os.fork()
    if pid == 0:
        os.close(read_fd)
        try:
            os.write(write_fd, b"1" if torch.cuda._is_in_bad_fork() else b"0")
        finally:
            os.close(write_fd)
            os._exit(0)
    os.close(write_fd)
    verdict = os.read(read_fd, 1)
    os.close(read_fd)
    _, status = os.waitpid(pid, 0)
    if status != 0 or verdict != b"0":
        raise RuntimeError(f"ZeroGPU parent CUDA fork poisoned during {event}")
    print(f'[WAN_DIAG] {{"bad_fork": false, "event": "{event}"}}', flush=True)


class WrapperLoopRuntime:
    _AOT_SHA256 = {
        "high": "58480b374bcdd47b50f2ea1351342bdb436828e298067a677368dcf53d473f97",
        "low": "f3ea9d67ec500d28e14c0ec61f77f52b812139bf440ded5dc27c211c184734a1",
    }

    def __init__(
        self,
        *,
        models_root: Path,
        high_model_name: str,
        low_model_name: str,
        clip_name: str,
        int8_clip_name: str,
        vae_name: str,
        sampler_name: str,
        scheduler_mode: str,
        split_step: int,
        riflex_k: int,
        loop_shift_skip: int,
        loop_start_percent: float,
        loop_end_percent: float,
        start_latent_strength: float,
        end_latent_strength: float,
        end_temporal_mask_strength: float,
        decode_end_image_hint: bool,
        fun_or_fl2v_model: bool,
        zero_end_latent_conditioning: bool,
        end_latent_conditioning_strength: float,
        low_pass_end_conditioning_strength: float,
        custom_sigmas: tuple[float, ...],
        attention_mode: str = "sdpa",
        text_encoder_quantization: str = "int8_weight_only",
        global_resident_models: bool = False,
        vae_tiling: bool = False,
    ) -> None:
        self.models_root = Path(models_root)
        self.high_model_name = high_model_name
        self.low_model_name = low_model_name
        self.clip_name = clip_name
        self.int8_clip_name = int8_clip_name
        self.vae_name = vae_name
        self.sampler_name = sampler_name
        self.scheduler_mode = scheduler_mode
        self.split_step = int(split_step)
        self.riflex_k = int(riflex_k)
        self.loop_shift_skip = int(loop_shift_skip)
        self.loop_start_percent = float(loop_start_percent)
        self.loop_end_percent = float(loop_end_percent)
        self.start_latent_strength = float(start_latent_strength)
        self.end_latent_strength = float(end_latent_strength)
        self.end_temporal_mask_strength = float(end_temporal_mask_strength)
        self.decode_end_image_hint = bool(decode_end_image_hint)
        self.fun_or_fl2v_model = bool(fun_or_fl2v_model)
        self.zero_end_latent_conditioning = bool(zero_end_latent_conditioning)
        self.end_latent_conditioning_strength = float(end_latent_conditioning_strength)
        self.low_pass_end_conditioning_strength = float(low_pass_end_conditioning_strength)
        self.custom_sigmas = tuple(float(x) for x in custom_sigmas)
        self.attention_mode = attention_mode
        self.text_encoder_quantization = text_encoder_quantization
        self.global_resident_models = bool(global_resident_models)
        self.vae_tiling = bool(vae_tiling)

        self._configure_model_paths()
        _assert_zerogpu_parent_fork_clean("parent.models_localized_clean")
        load_wrapper_module()
        _assert_zerogpu_parent_fork_clean("parent.wrapper_import_clean")

        self._comfy_nodes = importlib.import_module("nodes")
        self._model_loader = get_wrapper_node_class("WanVideoModelLoader")()
        self._vae_loader = get_wrapper_node_class("WanVideoVAELoader")()
        self._t5_loader = get_wrapper_node_class("LoadWanVideoT5TextEncoder")()
        self._text_encode = get_wrapper_node_class("WanVideoTextEncode")()
        self._image_encode = get_wrapper_node_class("WanVideoImageToVideoEncode")()
        self._sampler = get_wrapper_node_class("WanVideoSampler")()
        self._loop_args = get_wrapper_node_class("WanVideoLoopArgs")()
        self._decode = get_wrapper_node_class("WanVideoDecode")()
        # Diffusion weights are loaded just in time.  High and low must never
        # coexist on the ZeroGPU device.
        self.high_model = None
        self.low_model = None
        self.vae = None
        self.t5 = None
        self.clip = None
        self._text_mode = ""
        self._text_encoder_released = False
        self._text_encoder_on_cuda = False
        self._aot_package_cache: dict[str, str] = {}
        self._aot_prepared_roles: set[str] = set()
        self._load_text_encoder()
        _assert_zerogpu_parent_fork_clean("parent.text_encoder_clean")
        if self.global_resident_models:
            self._ensure_text_encoder_cuda()
            self._preload_global_models()
        self._diag("runtime.ready")

    def _aot_package_dir(self, role: str) -> str:
        cached = self._aot_package_cache.get(role)
        if cached is not None:
            return cached
        if os.environ.get("WAN_DISABLE_AOT", "0") == "1":
            self._diag("aot.disabled", role=role, reason="environment")
            self._aot_package_cache[role] = ""
            return ""
        package_dir = Path(__file__).resolve().parent.parent / "aot_artifacts" / role / "package"
        package_file = package_dir / "root" / "package.pt2"
        if not package_file.is_file():
            self._diag("aot.disabled", role=role, reason="package_missing")
            self._aot_package_cache[role] = ""
            return ""
        torch_base = torch.__version__.split("+", 1)[0]
        zerogpu_parent = False
        if os.environ.get("SPACES_ZERO_GPU", "").strip().lower() in {"1", "true", "yes", "on"}:
            try:
                from spaces.zero import wrappers as zero_wrappers

                zerogpu_parent = not zero_wrappers.forked
            except Exception:
                zerogpu_parent = True
        # The parent CUDA firewall must report is_available=False to keep the
        # later worker fork clean. ZeroGPU still exposes an sm_120 virtual
        # Blackwell device to global tensors, so select that pinned production
        # target without touching a real CUDA discovery API in the parent.
        if zerogpu_parent:
            capability = (12, 0)
        else:
            capability = torch.cuda.get_device_capability() if torch.cuda.is_available() else None
        if torch_base != "2.11.0" or capability != (12, 0):
            self._diag(
                "aot.disabled",
                role=role,
                reason="incompatible_runtime",
                torch_base=torch_base,
                capability=list(capability) if capability else None,
            )
            self._aot_package_cache[role] = ""
            return ""
        digest = hashlib.sha256(package_file.read_bytes()).hexdigest()
        if digest != self._AOT_SHA256[role]:
            raise RuntimeError(f"AOT package integrity check failed for role {role}")
        self._diag("aot.ready", role=role)
        self._aot_package_cache[role] = str(package_dir)
        return self._aot_package_cache[role]

    @staticmethod
    def _gpu_memory() -> dict[str, float | int | str]:
        # ZeroGPU forks its GPU worker. A real CUDA call in the parent poisons
        # that fork with "Cannot re-initialize CUDA". Identity APIs are
        # emulated, but mem_get_info/allocator APIs are not, so defer all of
        # them until spaces marks the child worker as forked.
        if os.environ.get("SPACES_ZERO_GPU", "").strip().lower() in {"1", "true", "yes", "on"}:
            try:
                from spaces.zero import wrappers as zero_wrappers

                if not zero_wrappers.forked:
                    return {"device": "zerogpu_parent", "cuda_initialized": False}
            except Exception:
                return {"device": "zerogpu_parent", "cuda_initialized": False}
        try:
            if not torch.cuda.is_available():
                return {"cuda": "unavailable"}
            free, total = torch.cuda.mem_get_info()
            gib = 1024**3
            return {
                "device": str(torch.cuda.get_device_name()),
                "allocated_gib": round(torch.cuda.memory_allocated() / gib, 3),
                "reserved_gib": round(torch.cuda.memory_reserved() / gib, 3),
                "peak_allocated_gib": round(torch.cuda.max_memory_allocated() / gib, 3),
                "free_gib": round(free / gib, 3),
                "total_gib": round(total / gib, 3),
            }
        except BaseException as exc:
            return {"cuda": f"stats_error:{type(exc).__name__}"}

    def _diag(self, event: str, **fields) -> None:
        payload = {
            "ts": datetime.now(timezone.utc).isoformat(timespec="milliseconds"),
            "event": event,
            **self._gpu_memory(),
            **fields,
        }
        print(f"[WAN_DIAG] {json.dumps(payload, sort_keys=True)}", flush=True)

    def _configure_model_paths(self) -> None:
        import sys

        comfy_runtime_root = (Path(__file__).resolve().parent / "third_party" / "comfy_runtime").resolve()
        if str(comfy_runtime_root) not in sys.path:
            sys.path.insert(0, str(comfy_runtime_root))

        import folder_paths

        for key, subdir in (
            ("diffusion_models", "diffusion_models"),
            ("vae", "vae"),
            ("text_encoders", "text_encoders"),
        ):
            folder_paths.add_model_folder_path(key, str(self.models_root / subdir))

    def _load_model(self, model_name: str, role: str):
        started = time.perf_counter()
        self._diag("model.load.start", role=role)
        quantization = "fp8_e4m3fn_scaled_fast" if role == "high" else "fp8_e4m3fn_fast"
        model = self._model_loader.loadmodel(
            model=model_name,
            base_precision="fp16_fast",
            load_device="main_device",
            quantization=quantization,
            attention_mode=self.attention_mode,
            block_swap_args=None,
            rms_norm_function="default",
            keep_loaded_models=self.global_resident_models,
            # Read into transient host staging, then let the wrapper assign the
            # model to main_device and delete the state dict before sampling.
            # This avoids holding both checkpoint and parameters on the 40 GB
            # GPU; it is loading staging, never runtime CPU offload.
            checkpoint_load_device="cpu",
        )[0]
        self._diag(
            "model.load.done",
            role=role,
            elapsed_s=round(time.perf_counter() - started, 3),
        )
        return model

    def _ensure_diffusion_model(self, role: str):
        attribute = "high_model" if role == "high" else "low_model"
        model = getattr(self, attribute)
        if model is None:
            name = self.high_model_name if role == "high" else self.low_model_name
            model = self._load_model(name, role)
            setattr(self, attribute, model)
        else:
            self._diag("model.reused", role=role)
        return model

    def _materialize_diffusion_model(self, role: str) -> None:
        """Assign external checkpoint weights and bind AOT before any request."""
        from ComfyUI_WanVideoWrapper.nodes_model_loading import load_weights_assign

        patcher = self._ensure_diffusion_model(role)
        wrapped = patcher.model
        state_dict = wrapped["sd"]
        if state_dict is not None:
            self._diag("model.materialize.start", role=role, device="cpu")
            assigned_parameters = load_weights_assign(
                wrapped.diffusion_model,
                state_dict,
                wrapped["weight_dtype"],
                wrapped["base_dtype"],
            )
            wrapped["sd"] = None
            del state_dict
            gc.collect()
            self._diag(
                "model.materialize.done",
                role=role,
                device="cpu",
                assigned_parameters=assigned_parameters,
            )

            transfer_started = time.perf_counter()
            self._diag("model.device_transfer.start", role=role, device="cuda")
            wrapped.diffusion_model.to(torch.device("cuda"))
            self._diag(
                "model.device_transfer.done",
                role=role,
                device="cuda",
                elapsed_s=round(time.perf_counter() - transfer_started, 3),
            )

        package_dir = self._aot_package_dir(role)
        if package_dir and role not in self._aot_prepared_roles:
            import spaces

            self._diag("aot.bind.start", role=role)
            spaces.aoti_load_from_package_dir(wrapped.diffusion_model.blocks, package_dir)
            self._aot_prepared_roles.add(role)
            self._diag("aot.bind.done", role=role)

    def _sampling_aot_package(self, role: str) -> str:
        if role in self._aot_prepared_roles:
            return ""
        return self._aot_package_dir(role)

    def _load_vae(self) -> None:
        if self.vae is None:
            started = time.perf_counter()
            self._diag("vae.load.start")
            self.vae = self._vae_loader.loadmodel(
                self.vae_name,
                "bf16",
                load_device="main_device",
            )[0]
            self._diag("vae.load.done", elapsed_s=round(time.perf_counter() - started, 3))

    def _preload_global_models(self) -> None:
        started = time.perf_counter()
        self._diag("global_preload.start", vae_tiling=self.vae_tiling)
        self._load_vae()
        self._materialize_diffusion_model("high")
        self._materialize_diffusion_model("low")
        self._diag(
            "global_preload.done",
            elapsed_s=round(time.perf_counter() - started, 3),
            vae_tiling=self.vae_tiling,
        )

    def resolve_workflow_size(self, image: Image.Image, target_width: int, target_height: int, divisible_by: int = 16) -> tuple[int, int]:
        src_w, src_h = image.size
        if src_w <= 0 or src_h <= 0:
            return int(target_width), int(target_height)
        ratio = min(float(target_width) / float(src_w), float(target_height) / float(src_h))
        width = max(int(round(src_w * ratio)), 1)
        height = max(int(round(src_h * ratio)), 1)
        if divisible_by > 1:
            width = max(width - (width % divisible_by), divisible_by)
            height = max(height - (height % divisible_by), divisible_by)
        return int(width), int(height)

    def _pil_to_comfy_image(self, image: Image.Image, width: int, height: int) -> torch.Tensor:
        resized = image.convert("RGB").resize((int(width), int(height)), Image.LANCZOS)
        array = np.asarray(resized, dtype=np.float32) / 255.0
        return torch.from_numpy(array).unsqueeze(0)

    def _load_text_encoder(self) -> None:
        if self.t5 is not None or self.clip is not None:
            return
        started = time.perf_counter()
        self._diag("text_encoder.load.start")
        self.t5 = None
        self.clip = None
        if self.text_encoder_quantization == "int8_weight_only":
            # This is the production path. Fail fast rather than silently
            # recreating the multi-minute BF16 encoder path on the first job.
            self.t5 = self._load_quanto_text_encoder_int8()
            self._text_mode = "wrapper_t5"
        else:
            try:
                self.t5 = self._t5_loader.loadmodel(
                    self.clip_name,
                    "bf16",
                    load_device="main_device",
                    quantization="disabled",
                )[0]
                self._text_mode = "wrapper_t5"
            except Exception as exc:
                print(f"[INFO] Loop wrapper T5 non disponibile per {self.clip_name}: {exc}")
                self.clip = self._comfy_nodes.CLIPLoader().load_clip(self.clip_name, "wan", "default")[0]
                self._text_mode = "comfy_clip"
            self._text_encoder_on_cuda = True
        self._text_encoder_released = False
        self._diag(
            "text_encoder.load.done",
            mode=self._text_mode,
            quantization=self.text_encoder_quantization,
            elapsed_s=round(time.perf_counter() - started, 3),
        )

    def _load_quanto_text_encoder_int8(self) -> dict:
        """Materialize the wrapper-native UMT5 directly from a QINT8 artifact."""
        from optimum.quanto.nn import QLinear
        from optimum.quanto.quantize import _quantize_submodule
        from safetensors import safe_open
        from safetensors.torch import load_file
        from ComfyUI_WanVideoWrapper.wanvideo.modules.t5 import T5EncoderModel

        import folder_paths
        _assert_zerogpu_parent_fork_clean("parent.quanto_import_clean")

        model_path = folder_paths.get_full_path_or_raise("text_encoders", self.int8_clip_name)
        tokenizer_path = (
            Path(__file__).resolve().parent
            / "third_party"
            / "ComfyUI-WanVideoWrapper"
            / "configs"
            / "T5_tokenizer"
        )
        with safe_open(model_path, framework="pt", device="cpu") as checkpoint:
            metadata = checkpoint.metadata() or {}
        if metadata.get("quantization_format") != "quanto":
            raise RuntimeError("The mounted INT8 text encoder is not a Quanto checkpoint.")
        encoded_map = metadata.get("quantization_map_base64", "")
        try:
            quantization_map = json.loads(base64.b64decode(encoded_map).decode("utf-8"))
        except Exception as exc:
            raise RuntimeError("Invalid Quanto quantization map in the INT8 text encoder.") from exc
        if len(quantization_map) != 192:
            raise RuntimeError(
                f"Unexpected UMT5 quantization map: {len(quantization_map)} modules, expected 192."
            )

        encoder = T5EncoderModel(
            text_len=512,
            dtype=torch.bfloat16,
            device=torch.device("cpu"),
            state_dict=None,
            tokenizer_path=str(tokenizer_path),
            quantization="disabled",
        )
        _assert_zerogpu_parent_fork_clean("parent.t5_meta_model_clean")
        started = time.perf_counter()
        self._diag(
            "text_encoder.int8.load.start",
            implementation="wrapper_t5_quanto",
            quantized_modules=len(quantization_map),
        )
        wrap_started = time.perf_counter()
        self._diag("text_encoder.int8.wrap.start")
        for name, module in list(encoder.model.named_modules()):
            qconfig = quantization_map.get(name)
            if qconfig is None:
                continue
            weights = None if qconfig["weights"] == "none" else qconfig["weights"]
            activations = None if qconfig["activations"] == "none" else qconfig["activations"]
            _quantize_submodule(
                encoder.model,
                name,
                module,
                weights=weights,
                activations=activations,
            )
        self._diag(
            "text_encoder.int8.wrap.done",
            elapsed_s=round(time.perf_counter() - wrap_started, 3),
        )
        _assert_zerogpu_parent_fork_clean("parent.quanto_wrap_clean")

        # Quanto's public requantize helper materializes the complete BF16 UMT5
        # before applying QINT8. Instead, rebuild each quantized module directly
        # from its serialized tensors on CPU. Peak host memory stays close to the
        # final 8.2 GiB model. The parent keeps this prepared CPU representation;
        # the forked ZeroGPU worker performs the one-time CUDA transfer.
        materialize_started = time.perf_counter()
        self._diag("text_encoder.int8.materialize.start", device="cpu")
        state_load_started = time.perf_counter()
        state_dict = load_file(model_path, device="cpu")
        _assert_zerogpu_parent_fork_clean("parent.t5_state_load_clean")
        self._diag(
            "text_encoder.int8.state_load.done",
            device="cpu",
            tensors=len(state_dict),
            elapsed_s=round(time.perf_counter() - state_load_started, 3),
        )
        handled_keys: set[str] = set()
        quantized_linears = 0
        dequantized_embeddings = 0
        checkpoint_keys = tuple(state_dict)
        for module_index, module_name in enumerate(quantization_map, start=1):
            module = encoder.model.get_submodule(module_name)
            prefix = f"{module_name}."
            module_keys = [key for key in checkpoint_keys if key.startswith(prefix)]
            if isinstance(module, QLinear):
                for name, parameter in list(module.named_parameters(recurse=False)):
                    if parameter is not None and parameter.is_meta:
                        setattr(
                            module,
                            name,
                            torch.nn.Parameter(
                                torch.empty_like(parameter, device="cpu"),
                                requires_grad=parameter.requires_grad,
                            ),
                        )
                for name, buffer in list(module.named_buffers(recurse=False)):
                    if buffer is not None and buffer.is_meta:
                        setattr(module, name, torch.empty_like(buffer, device="cpu"))
                module.load_state_dict(
                    {key[len(prefix):]: state_dict.pop(key) for key in module_keys},
                    strict=True,
                )
                quantized_linears += 1
            elif isinstance(module, torch.nn.Embedding):
                data = state_dict.pop(f"{prefix}weight._data")
                scale = state_dict.pop(f"{prefix}weight._scale")
                module.weight = torch.nn.Parameter(
                    (data * scale).to(torch.bfloat16),
                    requires_grad=False,
                )
                for key in module_keys:
                    state_dict.pop(key, None)
                dequantized_embeddings += 1
            else:
                raise RuntimeError(
                    f"Unsupported Quanto module in UMT5 checkpoint: {type(module).__name__}"
                )
            handled_keys.update(module_keys)
            if module_index % 24 == 0 or module_index == len(quantization_map):
                self._diag(
                    "text_encoder.int8.modules.progress",
                    completed=module_index,
                    total=len(quantization_map),
                )

        for key, value in list(state_dict.items()):
            parent_name, parameter_name = key.rsplit(".", 1)
            parent = encoder.model.get_submodule(parent_name)
            current = getattr(parent, parameter_name)
            if not isinstance(current, torch.nn.Parameter):
                raise RuntimeError(f"Unexpected non-parameter UMT5 tensor: {key}")
            setattr(
                parent,
                parameter_name,
                torch.nn.Parameter(value, requires_grad=current.requires_grad),
            )
            del state_dict[key]
        del state_dict
        self._diag("text_encoder.int8.cpu_model.done")
        _assert_zerogpu_parent_fork_clean("parent.t5_modules_clean")
        meta_parameters = [name for name, parameter in encoder.model.named_parameters() if parameter.is_meta]
        if meta_parameters:
            raise RuntimeError(f"UMT5 contains unmaterialized parameters: {meta_parameters[:5]}")
        self._diag(
            "text_encoder.int8.materialize.done",
            device="cpu",
            quantized_linears=quantized_linears,
            dequantized_embeddings=dequantized_embeddings,
            elapsed_s=round(time.perf_counter() - materialize_started, 3),
        )

        encoder.quantization = "quanto_int8"
        encoder.weights_prepared = True
        gc.collect()
        self._diag(
            "text_encoder.int8.load.done",
            implementation="wrapper_t5_quanto",
            device="cpu",
            elapsed_s=round(time.perf_counter() - started, 3),
        )
        return {
            "model": encoder,
            "dtype": torch.bfloat16,
            "name": self.int8_clip_name,
        }

    def _ensure_text_encoder_cuda(self) -> None:
        if self.text_encoder_quantization != "int8_weight_only" or self._text_encoder_on_cuda:
            return
        from optimum.quanto.nn import QLinear
        from optimum.quanto.tensor.weights.qbytes import WeightQBytesTensor

        encoder = self.t5["model"]
        transfer_started = time.perf_counter()
        self._diag("text_encoder.int8.device_transfer.start", device="cuda")
        device_to = torch.device("cuda")
        for module in encoder.model.modules():
            if isinstance(module, QLinear):
                weight = module.weight
                module.weight = torch.nn.Parameter(
                    WeightQBytesTensor(
                        weight.qtype,
                        weight.axis,
                        weight.size(),
                        weight.stride(),
                        weight._data.to(device_to),
                        weight._scale.to(device_to),
                        weight.activation_qtype,
                        requires_grad=False,
                    ),
                    requires_grad=False,
                )
                if module.bias is not None:
                    module.bias = torch.nn.Parameter(
                        module.bias.to(device_to),
                        requires_grad=module.bias.requires_grad,
                    )
                for name, buffer in list(module.named_buffers(recurse=False)):
                    if buffer is not None:
                        setattr(module, name, buffer.to(device_to))
                continue
            for name, parameter in list(module.named_parameters(recurse=False)):
                if parameter is not None:
                    setattr(
                        module,
                        name,
                        torch.nn.Parameter(
                            parameter.to(device_to),
                            requires_grad=parameter.requires_grad,
                        ),
                    )
            for name, buffer in list(module.named_buffers(recurse=False)):
                if buffer is not None:
                    setattr(module, name, buffer.to(device_to))
        self._diag(
            "text_encoder.int8.device_transfer.done",
            device="cuda",
            elapsed_s=round(time.perf_counter() - transfer_started, 3),
        )
        encoder.device = device_to
        self._text_encoder_on_cuda = True
        gc.collect()

    def _build_text_embeds_from_encoder(self, prompt: str, negative_prompt: str):
        if self._text_mode == "wrapper_t5":
            return self._text_encode.process(
                positive_prompt=prompt,
                negative_prompt=negative_prompt,
                t5=self.t5,
                force_offload=False,
                model_to_offload=None,
                use_disk_cache=False,
                device="gpu",
            )[0]

        positive = self._encode_comfy_wan_prompt(prompt)
        negative = self._encode_comfy_wan_prompt(negative_prompt or "")
        return {
            "prompt_embeds": positive,
            "negative_prompt_embeds": negative,
            "echoshot": False,
        }

    def get_text_embeds(self, prompt: str, negative_prompt: str):
        if self._text_encoder_released or ((self.t5 is None) and (self.clip is None)):
            self._load_text_encoder()
        self._ensure_text_encoder_cuda()
        return self._build_text_embeds_from_encoder(
            (prompt or "").strip(),
            (negative_prompt or "").strip(),
        )

    def release_text_encoder(self, *, force: bool = False) -> None:
        if not force and self.text_encoder_quantization == "int8_weight_only":
            self._diag("text_encoder.retained", mode=self._text_mode)
            return
        t5, clip = self.t5, self.clip
        self.t5 = None
        self.clip = None
        self._text_encoder_released = True
        self._text_encoder_on_cuda = False
        del t5, clip
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        self._diag("text_encoder.destroyed")

    def _release_diffusion_model(self, attribute: str, *, force: bool = False) -> None:
        role = "high" if attribute == "high_model" else "low"
        if self.global_resident_models and not force:
            self._diag("model.retained", role=role)
            return
        self._diag("model.destroy.start", role=role)
        model = getattr(self, attribute, None)
        setattr(self, attribute, None)
        if model is not None:
            # The wrapper registers patchers in Comfy's loaded-model list while
            # sampling. Remove that bookkeeping reference without invoking an
            # unload-to-CPU path, then destroy the GPU-resident object.
            try:
                import comfy.model_management as model_management

                patcher = getattr(model, "patcher", model)
                model_management.current_loaded_models[:] = [
                    loaded
                    for loaded in model_management.current_loaded_models
                    if loaded._model() is not patcher
                ]
            except Exception:
                pass
            # The sampler assigns all checkpoint tensors as real parameters.
            # Sever the patcher -> WanVideoModel -> transformer chain instead
            # of moving it to CPU; this makes the GPU allocation collectible.
            try:
                wrapped = model.model
                wrapped.diffusion_model = None
                wrapped.pipeline.clear()
                model.model = None
            except Exception:
                pass
            try:
                del patcher
            except UnboundLocalError:
                pass
            del model
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        self._diag("model.destroy.done", role=role)

    def _release_vae(self, *, force: bool = False) -> None:
        if self.global_resident_models and not force:
            self._diag("vae.retained", tiled=self.vae_tiling)
            return
        self._diag("vae.destroy.start", present=self.vae is not None)
        vae = self.vae
        self.vae = None
        if vae is not None:
            del vae
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        self._diag("vae.destroy.done")

    def cleanup_job(self) -> None:
        self._diag("runtime.job_cleanup.start")
        self._release_diffusion_model("high_model")
        self._release_diffusion_model("low_model")
        self._release_vae()
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        self._diag("runtime.job_cleanup.done")

    def shutdown(self) -> None:
        self._diag("runtime.shutdown.start")
        self._release_diffusion_model("high_model", force=True)
        self._release_diffusion_model("low_model", force=True)
        self._release_vae(force=True)
        self.release_text_encoder(force=True)
        self._diag("runtime.shutdown.done")

    def close(self) -> None:
        """Backward-compatible process shutdown alias."""
        self.shutdown()

    def _sigmas(self):
        scheduler_mode = (self.scheduler_mode or "").strip().lower()
        if scheduler_mode == "fixed":
            if not self.custom_sigmas:
                return None
            return torch.tensor(self.custom_sigmas, dtype=torch.float32)
        if scheduler_mode == "simple":
            return self._simple_sigmas()
        if scheduler_mode == "linear_quadratic":
            return self._linear_quadratic_sigmas()
        return None

    def _flow_sigmas(self, timesteps: int = 1000) -> torch.Tensor:
        t = torch.arange(1, int(timesteps) + 1, dtype=torch.float32) / float(timesteps)
        shift = float(self.model_sampling_shift) if hasattr(self, "model_sampling_shift") else 1.0
        return shift * t / (1 + (shift - 1) * t)

    def _simple_sigmas(self) -> torch.Tensor:
        steps = max(int(getattr(self, "_current_steps", 0)), 1)
        sigmas = self._flow_sigmas()
        stride = len(sigmas) / steps
        values = [float(sigmas[-(1 + int(x * stride))]) for x in range(steps)]
        values.append(0.0)
        return torch.tensor(values, dtype=torch.float32)

    def _linear_quadratic_sigmas(self) -> torch.Tensor:
        steps = max(int(getattr(self, "_current_steps", 0)), 1)
        if steps == 1:
            values = [1.0, 0.0]
        else:
            threshold_noise = 0.025
            linear_steps = steps // 2
            quadratic_steps = steps - linear_steps
            linear = [i * threshold_noise / linear_steps for i in range(linear_steps)] if linear_steps else []
            threshold_step_diff = linear_steps - threshold_noise * steps
            quadratic_coef = threshold_step_diff / (linear_steps * quadratic_steps ** 2) if linear_steps and quadratic_steps else 0.0
            linear_coef = threshold_noise / linear_steps - 2 * threshold_step_diff / (quadratic_steps ** 2) if linear_steps and quadratic_steps else 0.0
            const = quadratic_coef * (linear_steps ** 2)
            quadratic = [quadratic_coef * (i ** 2) + linear_coef * i + const for i in range(linear_steps, steps)]
            values = list(reversed(linear + quadratic))
            values.append(0.0)
        return torch.tensor(values, dtype=torch.float32)

    def _start_end_temporal_mask(self, num_frames: int, height: int, width: int) -> torch.Tensor | None:
        end_strength = max(0.0, min(1.0, float(self.end_temporal_mask_strength)))
        if end_strength >= 1.0:
            return None
        adjusted_frames = ((int(num_frames) - 1) // 4) * 4 + 1
        base_frames = adjusted_frames if self.fun_or_fl2v_model else adjusted_frames + 1
        mask = torch.zeros(base_frames, int(height), int(width), dtype=torch.float32)
        mask[0].fill_(1.0)
        mask[-1].fill_(end_strength)
        return mask

    def _zero_end_latent_conditioning(self, image_embeds: dict) -> dict:
        if not self.zero_end_latent_conditioning and self.end_latent_conditioning_strength == 1.0:
            return image_embeds
        latent = image_embeds.get("image_embeds")
        if not torch.is_tensor(latent) or latent.ndim < 2 or latent.shape[1] <= 0:
            return image_embeds
        image_embeds = dict(image_embeds)
        latent = latent.clone()
        strength = 0.0 if self.zero_end_latent_conditioning else self.end_latent_conditioning_strength
        latent[:, -1:] *= max(0.0, float(strength))
        image_embeds["image_embeds"] = latent
        return image_embeds

    def _low_pass_image_embeds(self, image_embeds: dict) -> dict:
        strength = max(0.0, float(self.low_pass_end_conditioning_strength))
        if strength == 1.0:
            return image_embeds
        latent = image_embeds.get("image_embeds")
        if not torch.is_tensor(latent) or latent.ndim < 2 or latent.shape[1] <= 0:
            return image_embeds
        low_embeds = dict(image_embeds)
        latent = latent.clone()
        latent[:, -1:] *= strength
        low_embeds["image_embeds"] = latent
        return low_embeds

    def _encode_comfy_wan_prompt(self, text: str):
        tokens = self.clip.tokenize(text or "")
        encoded = self.clip.encode_from_tokens(tokens, return_dict=True)
        cond = encoded["cond"]
        if cond.ndim != 3:
            raise ValueError(f"Formato inatteso text encoder Wan: shape={tuple(cond.shape)}")
        return [frame.detach() for frame in cond]

    def generate_segment_iter(
        self,
        *,
        prompt: str,
        negative_prompt: str,
        start_image: Image.Image,
        end_image: Image.Image,
        width: int,
        height: int,
        num_frames: int,
        steps: int,
        cfg: float,
        shift: float,
        seed: int,
        progress_callback: Callable[..., None] | None = None,
    ):
        if torch.cuda.is_available():
            torch.cuda.reset_peak_memory_stats()
            torch.cuda.synchronize()
        started = stage_started = time.perf_counter()
        metrics: dict[str, float] = {}
        self._diag(
            "generation.start",
            width=int(width),
            height=int(height),
            frames=int(num_frames),
            steps=int(steps),
            split_step=int(self.split_step),
        )

        if progress_callback is not None:
            progress_callback(0.08, desc="Encoding prompt…")
        yield {"stage": "Encoding prompt…"}
        self._diag("text_encode.start")
        text_embeds = self.get_text_embeds(prompt, negative_prompt)
        self._diag("text_encode.done")
        self.release_text_encoder()
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        metrics["text_encode_s"] = time.perf_counter() - stage_started
        stage_started = time.perf_counter()
        yield {"stage": "Loading image encoder…"}

        if progress_callback is not None:
            progress_callback(0.18, desc="Loading image encoder…")
        yield {"stage": "Loading image encoder…"}
        self._load_vae()
        if progress_callback is not None:
            progress_callback(0.22, desc="Encoding image…")
        yield {"stage": "Encoding image…"}
        self._diag("vae.encode.start")
        image_embeds = self._image_encode.process(
            width=int(width),
            height=int(height),
            num_frames=int(num_frames),
            noise_aug_strength=0.0,
            start_latent_strength=float(self.start_latent_strength),
            end_latent_strength=float(self.end_latent_strength),
            force_offload=False,
            vae=self.vae,
            start_image=self._pil_to_comfy_image(start_image, width, height),
            end_image=self._pil_to_comfy_image(end_image, width, height),
            fun_or_fl2v_model=self.fun_or_fl2v_model,
            temporal_mask=self._start_end_temporal_mask(num_frames, height, width),
            tiled_vae=self.vae_tiling,
            keep_output_on_device=True,
        )[0]
        self._diag("vae.encode.done")
        image_embeds = self._zero_end_latent_conditioning(image_embeds)
        # The wrapper carries the VAE inside image_embeds for optional advanced
        # sampler paths. This simple I2V path does not use it while sampling.
        # Remove the reference and destroy the GPU VAE until final decode.
        image_embeds = dict(image_embeds)
        image_embeds["vae"] = None
        self._release_vae()
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        metrics["image_encode_s"] = time.perf_counter() - stage_started
        stage_started = time.perf_counter()
        yield {"stage": "Loading motion engine — pass 1 of 2…"}

        self._current_steps = int(steps)
        self.model_sampling_shift = float(shift)
        sigmas = self._sigmas()
        loop_args = None
        if self.loop_shift_skip > 0:
            loop_args = self._loop_args.process(
                shift_skip=int(self.loop_shift_skip),
                start_percent=float(self.loop_start_percent),
                end_percent=float(self.loop_end_percent),
            )[0]
        if progress_callback is not None:
            progress_callback(0.30, desc="Loading motion model — pass 1 of 2…")
        yield {"stage": "Loading motion engine — pass 1 of 2…"}
        self.high_model = self._ensure_diffusion_model("high")
        if progress_callback is not None:
            progress_callback(0.40, desc="Sampling motion — pass 1 of 2…")
        yield {"stage": "Sampling motion — pass 1 of 2…"}
        self._diag("sampling.start", role="high", start_step=0, end_step=int(self.split_step))
        high_samples, _ = self._sampler.process(
            model=self.high_model,
            image_embeds=image_embeds,
            text_embeds=text_embeds,
            steps=int(steps),
            cfg=float(cfg),
            shift=float(shift),
            seed=int(seed),
            force_offload=False,
            scheduler=self.sampler_name,
            riflex_freq_index=int(self.riflex_k),
            batched_cfg=False,
            rope_function="comfy",
            sigmas=sigmas,
            end_step=int(self.split_step),
            keep_loaded_models=True,
            release_state_dict=True,
            reset_peak_memory_stats=False,
            keep_output_on_device=True,
            aot_package_dir=self._sampling_aot_package("high"),
        )
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        self._diag(
            "sampling.done",
            role="high",
            elapsed_s=round(time.perf_counter() - stage_started, 3),
        )
        metrics["high_sampling_s"] = time.perf_counter() - stage_started
        stage_started = time.perf_counter()
        self._release_diffusion_model("high_model")
        yield {"stage": "Loading motion engine — pass 2 of 2…"}

        if progress_callback is not None:
            progress_callback(0.54, desc="Loading motion model — pass 2 of 2…")
        yield {"stage": "Loading motion engine — pass 2 of 2…"}
        self.low_model = self._ensure_diffusion_model("low")
        if progress_callback is not None:
            progress_callback(0.64, desc="Sampling motion — pass 2 of 2…")
        yield {"stage": "Sampling motion — pass 2 of 2…"}
        self._diag(
            "sampling.start",
            role="low",
            start_step=int(self.split_step),
            end_step=int(steps),
        )
        low_samples, _ = self._sampler.process(
            model=self.low_model,
            image_embeds=self._low_pass_image_embeds(image_embeds),
            text_embeds=text_embeds,
            samples=high_samples,
            steps=int(steps),
            cfg=float(cfg),
            shift=float(shift),
            seed=int(seed),
            force_offload=False,
            scheduler=self.sampler_name,
            riflex_freq_index=int(self.riflex_k),
            batched_cfg=False,
            rope_function="comfy",
            sigmas=sigmas,
            loop_args=loop_args,
            start_step=int(self.split_step),
            keep_loaded_models=True,
            release_state_dict=True,
            reset_peak_memory_stats=False,
            keep_output_on_device=True,
            aot_package_dir=self._sampling_aot_package("low"),
        )
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        self._diag(
            "sampling.done",
            role="low",
            elapsed_s=round(time.perf_counter() - stage_started, 3),
        )
        metrics["low_sampling_s"] = time.perf_counter() - stage_started
        stage_started = time.perf_counter()
        self._release_diffusion_model("low_model")
        yield {"stage": "Reloading image decoder…"}
        if not self.decode_end_image_hint:
            low_samples = dict(low_samples)
            low_samples["end_image"] = None
            low_samples["drop_last"] = True
        if progress_callback is not None:
            progress_callback(0.74, desc="Reloading image decoder…")
        yield {"stage": "Reloading image decoder…"}
        self._load_vae()
        if progress_callback is not None:
            progress_callback(0.78, desc="Decoding frames…")
        yield {"stage": "Decoding frames…"}
        self._diag("vae.decode.start")
        images = self._decode.decode(
            vae=self.vae,
            samples=low_samples,
            enable_vae_tiling=self.vae_tiling,
            tile_x=272,
            tile_y=272,
            tile_stride_x=144,
            tile_stride_y=128,
            normalization="default",
            keep_model_loaded=True,
            keep_output_on_device=True,
        )[0]
        self._diag("vae.decode.done")
        self._release_vae()
        if torch.cuda.is_available():
            torch.cuda.synchronize()
            metrics["peak_vram_gib"] = torch.cuda.max_memory_allocated() / (1024**3)
        metrics["decode_s"] = time.perf_counter() - stage_started
        metrics["wan_total_s"] = time.perf_counter() - started
        self.last_metrics = metrics
        self._diag("generation.done", **{key: round(value, 3) for key, value in metrics.items()})
        if progress_callback is not None:
            progress_callback(0.86, desc="Frames decoded")
        yield images