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