aduc_framework/managers/wan_manager.py

# aduc_framework/managers/wan_manager.py
# WanManager v0.1.4 (final)

import os
import platform
import shutil
import subprocess
import tempfile
import random
from typing import List, Any, Optional, Tuple

import numpy as np
import torch
from PIL import Image

# SDPA / FlashAttention context (PyTorch 2.1+ / 2.0 fallback)
try:
    from torch.nn.attention import sdpa_kernel, SDPBackend  # PyTorch 2.1+
    _SDPA_NEW = True
except Exception:
    from torch.backends.cuda import sdp_kernel as _legacy_sdp  # PyTorch 2.0
    _SDPA_NEW = False

from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
from diffusers.utils.export_utils import export_to_video


class WanManager:
    """
    Wan i2v Manager:
    - Banner com verificações PyTorch/CUDA/SDPA/GPUs no startup
    - 2 Transformers 3D (alto/baixo ruído), bf16, device_map='auto', max_memory por GPU
    - LoRA Lightning fundida e descarregada
    - SDPA com preferência por FlashAttention + fallback (efficient/math)
    - 3 batentes: image(t=0, peso 1), handle(k da UI alinhado a 1 (mod 4)), last(t final)
    - Fallback se a pipeline não suportar args customizados (handle/anchor)
    """

    MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
    TRANSFORMER_ID = "cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers"

    # Dimensões/frames
    MAX_DIMENSION = 832
    MIN_DIMENSION = 480
    DIMENSION_MULTIPLE = 16
    SQUARE_SIZE = 480

    FIXED_FPS = 16
    MIN_FRAMES_MODEL = 8
    MAX_FRAMES_MODEL = 81

    default_negative_prompt = (
        "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，"
        "JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，"
        "手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走,过曝，"
    )

    def __init__(self) -> None:
        # Banner de verificação
        self._print_env_banner()

        print("Loading models into memory. This may take a few minutes...")

        # Sharding automático com chaves válidas (inteiros e "cpu")
        n_gpus = torch.cuda.device_count()
        max_memory = {i: "45GiB" for i in range(n_gpus)}  # ajuste conforme VRAM
        max_memory["cpu"] = "120GiB"

        transformer = WanTransformer3DModel.from_pretrained(
            self.TRANSFORMER_ID,
            subfolder="transformer",
            torch_dtype=torch.bfloat16,
            device_map="auto",
            max_memory=max_memory,
        )
        transformer_2 = WanTransformer3DModel.from_pretrained(
            self.TRANSFORMER_ID,
            subfolder="transformer_2",
            torch_dtype=torch.bfloat16,
            device_map="auto",
            max_memory=max_memory,
        )

        self.pipe = WanImageToVideoPipeline.from_pretrained(
            self.MODEL_ID,
            transformer=transformer,
            transformer_2=transformer_2,
            torch_dtype=torch.bfloat16,
        )

        # Scheduler FlowMatch Euler (shift=32.0)
        self.pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
            self.pipe.scheduler.config, shift=32.0
        )

        # LoRA Lightning (fusão)
        print("Applying 8-step Lightning LoRA...")
        try:
            self.pipe.load_lora_weights(
                "Kijai/WanVideo_comfy",
                weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
                adapter_name="lightx2v",
            )
            self.pipe.load_lora_weights(
                "Kijai/WanVideo_comfy",
                weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
                adapter_name="lightx2v_2",
                load_into_transformer_2=True,
            )
            self.pipe.set_adapters(["lightx2v", "lightx2v_2"], adapter_weights=[1.0, 1.0])

            print("Fusing LoRA weights into the main model...")
            self.pipe.fuse_lora(adapter_names=["lightx2v"], lora_scale=3.0, components=["transformer"])
            self.pipe.fuse_lora(adapter_names=["lightx2v_2"], lora_scale=1.0, components=["transformer_2"])
            self.pipe.unload_lora_weights()
            print("Lightning LoRA successfully fused. Model is ready for fast 8-step generation.")
        except Exception as e:
            print(f"[WanManager] AVISO: Falha ao fundir LoRA Lightning (seguirá sem fusão): {e}")

        print("All models loaded. Service is ready.")

    # ---------- Banner/Checks ----------
    def _print_env_banner(self) -> None:
        def _safe_get(fn, default="n/a"):
            try:
                return fn()
            except Exception:
                return default

        torch_ver = getattr(torch, "__version__", "unknown")
        cuda_rt = getattr(torch.version, "cuda", "unknown")
        cudnn_ver = _safe_get(lambda: torch.backends.cudnn.version())
        cuda_ok = torch.cuda.is_available()
        n_gpu = torch.cuda.device_count() if cuda_ok else 0
        devs, total_vram, caps = [], [], []
        if cuda_ok:
            for i in range(n_gpu):
                props = torch.cuda.get_device_properties(i)
                devs.append(f"cuda:{i} {props.name}")
                total_vram.append(f"{props.total_memory/1024**3:.1f}GiB")
                caps.append(f"{props.major}.{props.minor}")

        # BF16/TF32
        try:
            bf16_supported = bool(getattr(torch.cuda, "is_bf16_supported", lambda: False)())
        except Exception:
            bf16_supported = False
            if cuda_ok and caps:
                major = int(caps[0].split(".")[0])
                bf16_supported = major >= 8
        tf32_allowed = getattr(torch.backends.cuda.matmul, "allow_tf32", False)

        # SDPA API
        try:
            from torch.nn.attention import sdpa_kernel as _probe1  # noqa
            sdpa_api = "torch.nn.attention (2.1+)"
        except Exception:
            try:
                from torch.backends.cuda import sdp_kernel as _probe2  # noqa
                sdpa_api = "torch.backends.cuda (2.0)"
            except Exception:
                sdpa_api = "unavailable"

        # xFormers
        try:
            import xformers  # noqa
            xformers_ok = True
        except Exception:
            xformers_ok = False

        alloc_conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "unset")
        visible = os.environ.get("CUDA_VISIBLE_DEVICES", "unset")
        python_ver = platform.python_version()
        nvcc = shutil.which("nvcc")
        nvcc_ver = "n/a"
        if nvcc:
            try:
                nvcc_ver = subprocess.check_output([nvcc, "--version"], text=True).strip().splitlines()[-1]
            except Exception:
                nvcc_ver = "n/a"

        banner_lines = [
            "================== WAN MANAGER • ENV ==================",
            f"Python              : {python_ver}",
            f"PyTorch             : {torch_ver}",
            f"CUDA (torch)        : {cuda_rt}",
            f"cuDNN               : {cudnn_ver}",
            f"CUDA available      : {cuda_ok}",
            f"GPU count           : {n_gpu}",
            f"GPUs                : {', '.join(devs) if devs else 'n/a'}",
            f"GPU VRAM            : {', '.join(total_vram) if total_vram else 'n/a'}",
            f"Compute Capability  : {', '.join(caps) if caps else 'n/a'}",
            f"BF16 supported      : {bf16_supported}",
            f"TF32 allowed        : {tf32_allowed}",
            f"SDPA API            : {sdpa_api}",
            f"xFormers available  : {xformers_ok}",
            f"CUDA_VISIBLE_DEVICES: {visible}",
            f"PYTORCH_CUDA_ALLOC_CONF: {alloc_conf}",
            f"nvcc                : {nvcc_ver}",
            "=======================================================",
        ]
        print("\n".join(banner_lines))

    # ---------- utils de imagem ----------
    def _round_multiple(self, x: int, multiple: int) -> int:
        return int(round(x / multiple) * multiple)

    def process_image_for_video(self, image: Image.Image) -> Image.Image:
        w, h = image.size
        if w == h:
            return image.resize((self.SQUARE_SIZE, self.SQUARE_SIZE), Image.Resampling.LANCZOS)

        ar = w / h
        nw, nh = w, h

        # clamp superior
        if nw > self.MAX_DIMENSION or nh > self.MAX_DIMENSION:
            s = (self.MAX_DIMENSION / nw) if ar > 1 else (self.MAX_DIMENSION / nh)
            nw, nh = nw * s, nh * s

        # clamp inferior
        if nw < self.MIN_DIMENSION or nh < self.MIN_DIMENSION:
            s = (self.MIN_DIMENSION / nh) if ar > 1 else (self.MIN_DIMENSION / nw)
            nw, nh = nw * s, nh * s

        fw = self._round_multiple(int(nw), self.DIMENSION_MULTIPLE)
        fh = self._round_multiple(int(nh), self.DIMENSION_MULTIPLE)

        # mínimos finais coerentes
        fw = max(fw, self.MIN_DIMENSION if ar < 1 else self.SQUARE_SIZE)
        fh = max(fh, self.MIN_DIMENSION if ar > 1 else self.SQUARE_SIZE)

        return image.resize((fw, fh), Image.Resampling.LANCZOS)

    def resize_and_crop_to_match(self, target: Image.Image, ref: Image.Image) -> Image.Image:
        rw, rh = ref.size
        tw, th = target.size
        s = max(rw / tw, rh / th)
        nw, nh = int(tw * s), int(th * s)
        resized = target.resize((nw, nh), Image.Resampling.LANCZOS)
        left, top = (nw - rw) // 2, (nh - rh) // 2
        return resized.crop((left, top, left + rw, top + rh))

    # ---------- API ----------
    def generate_video_from_conditions(
        self,
        images_condition_items: List[List[Any]],  # [[image(Image), frame(int|str), peso(float)], ...]
        prompt: str,
        negative_prompt: Optional[str],
        duration_seconds: float,
        steps: int,
        guidance_scale: float,
        guidance_scale_2: float,
        seed: int,
        randomize_seed: bool,
        output_type: str = "np",
    ) -> Tuple[str, int]:
        # validação
        if not images_condition_items or len(images_condition_items) < 2:
            raise ValueError("Forneça ao menos dois itens (início e fim).")

        items = images_condition_items
        start_image = items[0][0]
        end_image = items[-1][0]
        if start_image is None or end_image is None:
            raise ValueError("As imagens inicial e final não podem ser vazias.")
        if not isinstance(start_image, Image.Image) or not isinstance(end_image, Image.Image):
            raise TypeError("Patches devem ser PIL.Image.")

        # handle opcional
        handle_image = items[1][0] if len(items) >= 3 else None

        # pesos
        handle_weight = float(items[1][2]) if len(items) >= 3 and items[1][2] is not None else 1.0
        end_weight = float(items[-1][2]) if len(items[-1]) >= 3 and items[-1][2] is not None else 1.0

        # preprocess e alinhamento HxW
        processed_start = self.process_image_for_video(start_image)
        processed_end = self.resize_and_crop_to_match(end_image, processed_start)
        processed_handle = self.resize_and_crop_to_match(handle_image, processed_start) if handle_image else None

        H, W = processed_start.height, processed_start.width

        # frames (pipeline ajusta para 4n+1 internamente, aqui só clamp)
        num_frames = int(round(duration_seconds * self.FIXED_FPS))
        num_frames = int(np.clip(num_frames, self.MIN_FRAMES_MODEL, self.MAX_FRAMES_MODEL))

        # seed
        current_seed = random.randint(0, np.iinfo(np.int32).max) if randomize_seed else int(seed)
        generator = torch.Generator().manual_seed(current_seed)

        # argumentos base
        call_kwargs = dict(
            image=processed_start,
            last_image=processed_end,
            prompt=prompt,
            negative_prompt=negative_prompt if negative_prompt else self.default_negative_prompt,
            height=H,
            width=W,
            num_frames=num_frames,
            guidance_scale=float(guidance_scale),
            guidance_scale_2=float(guidance_scale_2),
            num_inference_steps=int(steps),
            generator=generator,
            output_type=output_type,
        )

        # mapear frame da UI do handle → índice latente alinhado a 1 (mod 4)
        corrected_handle_index = int(items[1][1])
        
        # Montar kwargs finais (com/sem handle)
        if processed_handle is not None:
            kwargs = dict(
                **call_kwargs,
                handle_image=processed_handle,
                handle_weight=float(handle_weight),
                handle_latent_index=corrected_handle_index,
                anchor_weight_last=float(end_weight),
            )
        else:
            kwargs = dict(
                **call_kwargs,
                anchor_weight_last=float(end_weight),
            )

        # Execução com SDPA e fallback de backend
        result = None
        
        result = self.pipe(**kwargs)

        frames = result.frames[0]

        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
            video_path = tmp.name
        export_to_video(frames, video_path, fps=self.FIXED_FPS)

        return video_path, current_seed