"""Qwen 2512 Canny regeneration followed by masked Z-Image face repair.

This profile ports the two-stage architecture used by cebeuq/Synthid-Bypass:

1. Qwen-Image-2512 img2img with the 4-step Lightning LoRA and the DiffSynth
   blockwise Canny ControlNet regenerates the whole image, optionally as
   overlapping feather-blended tiles for large inputs.
2. Faces are detected on the original image, refined to masks with SAM, regenerated
   from the original face crops with Z-Image Turbo, and feathered into stage 1.

The runtime intentionally uses permissively licensed YuNet instead of the reference
workflow's Ultralytics detector. All diffusion and segmentation models remain the same
model families. The adaptive formulas are direct ports, while the face result is scaled
for this runtime's different sampler and mask-compositing path.
"""

# DiffSynth, torch, transformers, and cv2 expose mostly untyped tensor/array APIs.
# Keep the relaxation local to this optional ML boundary.
# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportUnnecessaryComparison=false
from __future__ import annotations

import contextlib
import hashlib
import logging
import math
import os
import tempfile
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any

import numpy as np
from PIL import Image

from remove_ai_watermarks.noai.watermark_profiles import resolve_seed

if TYPE_CHECKING:
    from collections.abc import Callable

log = logging.getLogger(__name__)

QWEN_IMAGE_2512_MODEL_ID = "Qwen/Qwen-Image-2512"
QWEN_CANNY_CONTROLNET_MODEL_ID = "DiffSynth-Studio/Qwen-Image-Blockwise-ControlNet-Canny"
QWEN_LIGHTNING_MODEL_ID = "lightx2v/Qwen-Image-2512-Lightning"
QWEN_LIGHTNING_PATTERN = "Qwen-Image-2512-Lightning-4steps-V1.0-bf16.safetensors"
ZIMAGE_TURBO_MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
SAM_MODEL_ID = "facebook/sam-vit-base"

YUNET_MODEL_URL = (
    "https://media.githubusercontent.com/media/opencv/opencv_zoo/main/"
    "models/face_detection_yunet/face_detection_yunet_2023mar.onnx"
)
YUNET_MODEL_NAME = "face_detection_yunet_2023mar.onnx"
YUNET_MODEL_SHA256 = "8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4"
# The upstream graph's 0.2 threshold belongs to YOLO and does not transfer to
# YuNet's score calibration. At 0.2 YuNet admitted background and decorative
# false positives, multiplying the serial Z-Image face-stage cost. A 0.5 gate
# retained all visible faces in the public and upstream comparison fixtures.
YUNET_SCORE_THRESHOLD = 0.5

GLOBAL_STEPS = 4
FACE_STEPS = 8
GLOBAL_CFG = 1.0
FACE_CFG = 1.0
GLOBAL_CONTROLNET_SCALE = 1.0
RESIDENT_FACE_MODEL_MIN_VRAM_GIB = 64.0
# The reference face denoise assumes its ComfyUI detailer sampler, latent
# noise-mask feather, and inpaint path. Applying that value unchanged to this
# DiffSynth crop-regeneration port over-processes faces. Paired public-fixture
# measurements and both provider oracles certified half the reference value.
FACE_DENOISE_SCALE = 0.5

# The source graph uses normalized Canny thresholds 0.05 and 0.25. OpenCV takes
# byte thresholds, so round 255*x to the matching integer values.
_CANNY_LOW = 13
_CANNY_HIGH = 64

# These strings intentionally preserve the reference workflow spelling. They are
# model inputs, not user-facing copy, and changing them would change the port.
_GLOBAL_PROMPT = "ultra clear and smoothe skin, spotless skin"
_GLOBAL_NEGATIVE = "moles, freckes, high detail skin"
_FACE_PROMPT = ""
_FACE_NEGATIVE = "blurry, ugly, bad quality,"


def resolve_face_model_residency(
    requested: bool | None,
    *,
    total_memory_gib: float,
) -> bool:
    """Keep the Z-Image stack resident when explicitly requested or safely sized."""
    if requested is not None:
        return requested
    return total_memory_gib >= RESIDENT_FACE_MODEL_MIN_VRAM_GIB


def _pin_vram_managed_models(pipe: Any) -> None:
    """Move the managed Z-Image stack to CUDA once and make offload a no-op."""
    model_names = ["text_encoder", "dit", "vae_encoder", "vae_decoder"]
    for name in model_names:
        model = getattr(pipe, name, None)
        if model is None:
            continue
        for module in model.modules():
            if not all(
                hasattr(module, attribute)
                for attribute in (
                    "offload_dtype",
                    "offload_device",
                    "onload_dtype",
                    "onload_device",
                    "preparing_dtype",
                    "preparing_device",
                    "computation_dtype",
                    "computation_device",
                )
            ):
                continue
            module.offload_dtype = module.computation_dtype
            module.offload_device = module.computation_device
            module.onload_dtype = module.computation_dtype
            module.onload_device = module.computation_device
            module.preparing_dtype = module.computation_dtype
            module.preparing_device = module.computation_device
    pipe.load_models_to_device(model_names)


def _cached_prompt_process(original_process: Any) -> Any:
    cache: dict[str, dict[str, Any]] = {}

    def cached_process(
        runtime_pipe: Any,
        prompt: str,
        edit_image: Any = None,
    ) -> dict[str, Any]:
        if edit_image is not None:
            return original_process(runtime_pipe, prompt, edit_image=edit_image)
        if prompt not in cache:
            cache[prompt] = original_process(runtime_pipe, prompt, edit_image=None)
        return cache[prompt]

    return cached_process


def _cache_static_prompt_embeddings(
    pipe: Any,
    output_params: tuple[str, ...],
) -> bool:
    """Memoize a prompt unit when its embedding depends only on static text."""
    for unit in pipe.units:
        if tuple(getattr(unit, "output_params", ())) == output_params:
            unit.process = _cached_prompt_process(unit.process)
            return True
    return False


def _clamp(value: float, minimum: float, maximum: float) -> float:
    if maximum < minimum:
        minimum, maximum = maximum, minimum
    return max(minimum, min(maximum, value))


def resolution_adaptive_denoise(
    width: int,
    height: int,
    *,
    adaptive_level: int = 6,
    denoise_min: float = 0.08,
    denoise_max: float = 0.15,
) -> float:
    """Port the reference resolution-based adaptive denoise calculation.

    At neutral level 5, 0.30 MP maps to ``denoise_min`` and 3.70 MP maps to
    ``denoise_max``. Levels above or below 5 add the same asymmetric spread as the
    reference custom node.
    """
    image_mp = max(1.0, float(width) * float(height)) / 1_000_000.0
    normalized = _clamp((image_mp - 0.30) / (3.70 - 0.30), 0.0, 1.0)

    minimum = float(denoise_min)
    maximum = float(denoise_max)
    if maximum < minimum:
        minimum, maximum = maximum, minimum
    denoise_range = maximum - minimum
    base = minimum + denoise_range * normalized

    level = int(adaptive_level)
    if level >= 5:
        offset = ((float(level) - 5.0) / 5.0) * denoise_range * 0.285714
    else:
        offset = -((5.0 - float(level)) / 4.0) * denoise_range * 0.257143
    return _clamp(base + offset, 0.0001, 1.0)


def largest_face_denoise(
    boxes: list[tuple[int, int, int, int]],
    image_size: tuple[int, int],
    *,
    base_denoise: float = 0.10,
    adaptive_ratio: float = 0.03,
    denoise_min: float = 0.05,
    denoise_max: float = 0.28,
) -> float:
    """Scale face denoise from the largest face area, matching reference mode."""
    width, height = image_size
    image_area = max(1.0, float(width) * float(height))
    largest_ratio = 0.0
    for x1, y1, x2, y2 in boxes:
        box_area = max(0.0, float(x2 - x1)) * max(0.0, float(y2 - y1))
        largest_ratio = max(largest_ratio, box_area / image_area)
    if largest_ratio <= 0.0:
        return _clamp(base_denoise, denoise_min, denoise_max)
    scaled = float(base_denoise) * (largest_ratio / max(1e-6, float(adaptive_ratio)))
    return _clamp(scaled, denoise_min, denoise_max)


def _target_size(width: int, height: int) -> tuple[int, int]:
    """Floor image dimensions to the /16 latent grid without changing aspect."""
    return max(16, (width // 16) * 16), max(16, (height // 16) * 16)


def _resize_to_target(image: Image.Image) -> Image.Image:
    """Resize pixels to the exact latent-grid dimensions passed to DiffSynth."""
    target = _target_size(image.width, image.height)
    if image.size == target:
        return image
    return image.resize(target, Image.Resampling.LANCZOS)


def build_canny_control_image(image: Image.Image) -> Image.Image:
    """Build the three-channel Canny conditioning image used by stage 1."""
    import cv2

    rgb = np.asarray(image.convert("RGB"))
    gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
    edges = cv2.Canny(gray, _CANNY_LOW, _CANNY_HIGH)
    return Image.fromarray(np.repeat(edges[:, :, None], 3, axis=2))


def build_global_kwargs(
    image: Image.Image,
    *,
    strength: float,
    seed: int | None,
    controlnet_input: Any,
) -> dict[str, Any]:
    """Build the DiffSynth Qwen call shape without importing the ML runtime."""
    input_image = _resize_to_target(image)
    width, height = input_image.size
    return {
        "prompt": _GLOBAL_PROMPT,
        "negative_prompt": _GLOBAL_NEGATIVE,
        "cfg_scale": GLOBAL_CFG,
        "input_image": input_image,
        "denoising_strength": float(strength),
        "height": height,
        "width": width,
        "seed": seed,
        "rand_device": "cpu",
        "num_inference_steps": GLOBAL_STEPS,
        # The source graph applies ModelSamplingAuraFlow with shift=3. DiffSynth
        # expresses the same rational sigma shift as exp(mu), so mu=log(3).
        "exponential_shift_mu": math.log(3.0),
        "blockwise_controlnet_inputs": [controlnet_input],
    }


def build_face_kwargs(crop: Image.Image, *, strength: float, seed: int | None) -> dict[str, Any]:
    """Build the DiffSynth Z-Image face-detail call shape."""
    input_image = _resize_to_target(crop)
    width, height = input_image.size
    return {
        "prompt": _FACE_PROMPT,
        "negative_prompt": _FACE_NEGATIVE,
        "cfg_scale": FACE_CFG,
        "input_image": input_image,
        "denoising_strength": float(strength),
        "height": height,
        "width": width,
        "seed": seed,
        "rand_device": "cpu",
        "num_inference_steps": FACE_STEPS,
    }


def composite_face(base: np.ndarray, detail: np.ndarray, mask: np.ndarray, *, feather: int = 10) -> np.ndarray:
    """Feather ``detail`` into ``base`` while preserving every zero-mask pixel."""
    import cv2

    if base.shape != detail.shape:
        raise ValueError("base and detail must have identical shapes")
    if mask.shape != base.shape[:2]:
        raise ValueError("mask must match the image height and width")

    alpha = mask.astype(np.float32) / 255.0
    if feather > 0:
        sigma = max(0.1, float(feather) / 3.0)
        alpha = cv2.GaussianBlur(alpha, (0, 0), sigmaX=sigma, sigmaY=sigma)
        # Blurring may introduce tiny values far outside the intended mask. Keep an
        # explicit support dilation so pixels beyond the feather radius stay exact.
        support = cv2.dilate((mask > 0).astype(np.uint8), np.ones((2 * feather + 1, 2 * feather + 1), np.uint8))
        alpha *= support
    alpha = np.clip(alpha, 0.0, 1.0)[:, :, None]
    merged = base.astype(np.float32) * (1.0 - alpha) + detail.astype(np.float32) * alpha
    return np.clip(np.rint(merged), 0, 255).astype(np.uint8)


def _expanded_box(
    box: tuple[int, int, int, int],
    image_size: tuple[int, int],
    *,
    factor: float = 2.5,
) -> tuple[int, int, int, int]:
    """Expand a face box around its center, matching the reference crop factor."""
    x1, y1, x2, y2 = box
    image_width, image_height = image_size
    center_x = (x1 + x2) / 2.0
    center_y = (y1 + y2) / 2.0
    width = max(1.0, (x2 - x1) * factor)
    height = max(1.0, (y2 - y1) * factor)
    return (
        max(0, round(center_x - width / 2.0)),
        max(0, round(center_y - height / 2.0)),
        min(image_width, round(center_x + width / 2.0)),
        min(image_height, round(center_y + height / 2.0)),
    )


def _model_cache_dir() -> Path:
    root = os.environ.get("XDG_CACHE_HOME")
    base = Path(root) if root else Path.home() / ".cache"
    return base / "remove-ai-watermarks"


def _yunet_model_path() -> Path:
    """Download the small MIT-licensed YuNet ONNX model on first use."""
    model_path = _model_cache_dir() / YUNET_MODEL_NAME
    if model_path.exists() and hashlib.sha256(model_path.read_bytes()).hexdigest() == YUNET_MODEL_SHA256:
        return model_path

    model_path.parent.mkdir(parents=True, exist_ok=True)
    log.info("Downloading YuNet face detector: %s", YUNET_MODEL_URL)
    request = urllib.request.Request(
        YUNET_MODEL_URL,
        headers={"User-Agent": "remove-ai-watermarks"},
    )
    with urllib.request.urlopen(request, timeout=60) as response:  # noqa: S310 - fixed HTTPS source
        payload = response.read()
        log.info(
            "YuNet download response: status=%s content_length=%s",
            getattr(response, "status", None),
            len(payload),
        )
    digest = hashlib.sha256(payload).hexdigest()
    if digest != YUNET_MODEL_SHA256:
        raise OSError(
            "YuNet download failed integrity verification: "
            f"expected {YUNET_MODEL_SHA256}, got {digest} ({len(payload)} bytes)"
        )
    with tempfile.NamedTemporaryFile(dir=model_path.parent, suffix=".onnx", delete=False) as handle:
        handle.write(payload)
        temporary = Path(handle.name)
    temporary.replace(model_path)
    return model_path


def _nms_boxes(
    boxes: list[tuple[int, int, int, int]],
    scores: list[float],
    *,
    threshold: float = 0.3,
) -> list[tuple[int, int, int, int]]:
    """Apply OpenCV NMS to detections collected at multiple image scales."""
    if not boxes:
        return []
    import cv2

    xywh = [[x1, y1, x2 - x1, y2 - y1] for x1, y1, x2, y2 in boxes]
    indices = cv2.dnn.NMSBoxes(xywh, scores, score_threshold=0.2, nms_threshold=threshold)
    if len(indices) == 0:
        return []
    return [boxes[int(index)] for index in np.asarray(indices).reshape(-1)]


def detect_faces(image: Image.Image) -> list[tuple[int, int, int, int]]:
    """Detect face boxes with YuNet at two scales for large and small faces."""
    import cv2

    rgb = np.asarray(image.convert("RGB"))
    original_height, original_width = rgb.shape[:2]
    detections: list[tuple[int, int, int, int]] = []
    scores: list[float] = []
    model_path = _yunet_model_path()

    for long_side in (640, 1280):
        scale = min(1.0, long_side / max(original_width, original_height))
        width = max(1, round(original_width * scale))
        height = max(1, round(original_height * scale))
        resized = cv2.resize(rgb, (width, height), interpolation=cv2.INTER_AREA) if scale < 1.0 else rgb
        bgr = cv2.cvtColor(resized, cv2.COLOR_RGB2BGR)
        detector = cv2.FaceDetectorYN.create(
            str(model_path),
            "",
            (width, height),
            YUNET_SCORE_THRESHOLD,
            0.3,
            5000,
        )
        _, rows = detector.detect(bgr)
        if rows is None:
            continue
        inverse = 1.0 / scale
        for row in rows:
            x, y, box_width, box_height = (float(value) for value in row[:4])
            x1 = max(0, round(x * inverse))
            y1 = max(0, round(y * inverse))
            x2 = min(original_width, round((x + box_width) * inverse))
            y2 = min(original_height, round((y + box_height) * inverse))
            if x2 <= x1 or y2 <= y1:
                continue
            detections.append((x1, y1, x2, y2))
            scores.append(float(row[-1]))
        if scale == 1.0:
            break
    return _nms_boxes(detections, scores)


def _ellipse_masks(
    boxes: list[tuple[int, int, int, int]],
    image_size: tuple[int, int],
) -> list[np.ndarray]:
    """Safe fallback masks when SAM is unavailable."""
    import cv2

    width, height = image_size
    masks: list[np.ndarray] = []
    for x1, y1, x2, y2 in boxes:
        mask = np.zeros((height, width), dtype=np.uint8)
        center = ((x1 + x2) // 2, (y1 + y2) // 2)
        axes = (max(1, int((x2 - x1) * 0.55)), max(1, int((y2 - y1) * 0.62)))
        cv2.ellipse(mask, center, axes, 0, 0, 360, 255, -1)
        masks.append(mask)
    return masks


def _prepare_sam_inputs(inputs: Any, device: str, dtype: Any) -> Any:
    """Move SAM inputs to the target device without casting geometric prompts."""
    prepared = inputs.to(device)
    if "pixel_values" in prepared:
        prepared["pixel_values"] = prepared["pixel_values"].to(dtype=dtype)
    return prepared


def _sam_point_prompts(
    boxes: list[tuple[int, int, int, int]],
) -> tuple[list[list[list[list[float]]]], list[list[list[int]]]]:
    """Build Impact Pack's center-1 positive prompt for every face box."""
    points = [[[[(x1 + x2) / 2.0, (y1 + y2) / 2.0]] for x1, y1, x2, y2 in boxes]]
    labels = [[[1] for _box in boxes]]
    return points, labels


def _clip_sam_masks_to_boxes(
    masks: list[np.ndarray],
    boxes: list[tuple[int, int, int, int]],
    image_size: tuple[int, int],
) -> list[np.ndarray]:
    """Match Impact Pack by intersecting each SAM mask with its detector box."""
    width, height = image_size
    clipped: list[np.ndarray] = []
    for mask, (x1, y1, x2, y2) in zip(masks, boxes, strict=True):
        box_mask = np.zeros((height, width), dtype=np.uint8)
        box_mask[max(0, y1) : min(height, y2), max(0, x1) : min(width, x2)] = 255
        clipped.append(np.bitwise_and(mask.astype(np.uint8), box_mask))
    return clipped


def _select_sam_masks(
    masks: np.ndarray,
    scores: np.ndarray,
    *,
    threshold: float = 0.93,
) -> list[np.ndarray]:
    """Select and combine SAM proposals like Impact Pack's ``sub_threshold``."""
    mask_array = np.asarray(masks)
    score_array = np.asarray(scores)
    if mask_array.ndim == 5 and mask_array.shape[0] == 1:
        mask_array = mask_array[0]
    if score_array.ndim == 3 and score_array.shape[0] == 1:
        score_array = score_array[0]
    if mask_array.ndim == 3:
        mask_array = mask_array[:, None, :, :]
    if score_array.ndim == 1:
        score_array = score_array[:, None]
    if mask_array.ndim != 4 or score_array.ndim != 2:
        raise ValueError("SAM masks and scores have unexpected dimensions")
    if mask_array.shape[:2] != score_array.shape:
        raise ValueError("SAM mask proposals and IoU scores do not align")

    selected_masks: list[np.ndarray] = []
    for candidates, candidate_scores in zip(mask_array, score_array, strict=True):
        selected = np.flatnonzero(candidate_scores >= threshold)
        if selected.size == 0:
            selected = np.asarray([int(np.argmax(candidate_scores))])
        combined = np.any(candidates[selected] > 0, axis=0)
        selected_masks.append(combined.astype(np.uint8) * 255)
    return selected_masks


def _sam_outputs_to_numpy(masks: Any, scores: Any) -> tuple[np.ndarray, np.ndarray]:
    """Convert SAM tensors through float32 because NumPy rejects bfloat16."""
    mask_array = masks.detach().float().cpu().numpy()
    score_array = scores.detach().float().cpu().numpy()
    return mask_array, score_array


@dataclass
class QwenZImagePipeline:
    """Lazy runtime for the two-stage Qwen/Z-Image profile."""

    device: str
    torch_dtype: Any
    hf_token: str | None = None
    progress_callback: Callable[[str], None] | None = None
    controlnet_conditioning_scale: float = GLOBAL_CONTROLNET_SCALE
    keep_face_models_on_device: bool | None = None
    cache_prompt_embeddings: bool = True

    def __post_init__(self) -> None:
        self._qwen_pipe: Any = None
        self._zimage_pipe: Any = None
        self._sam_model: Any = None
        self._sam_processor: Any = None

    def _progress(self, message: str) -> None:
        if self.progress_callback is not None:
            with contextlib.suppress(Exception):
                self.progress_callback(message)

    def _require_cuda(self) -> None:
        if self.device != "cuda":
            raise RuntimeError(
                "The qwen-zimage pipeline is CUDA-only. Its Qwen-Image-2512 and "
                "Z-Image models do not fit the supported MPS path."
            )

    def _vram_limit(self) -> float | None:
        import torch

        with contextlib.suppress(Exception):
            return max(1.0, torch.cuda.mem_get_info("cuda")[1] / (1024**3) - 0.5)
        return None

    def _keep_face_models_resident(self) -> bool:
        import torch

        total_memory_gib = 0.0
        with contextlib.suppress(Exception):
            total_memory_gib = torch.cuda.get_device_properties("cuda").total_memory / (1024**3)
        return resolve_face_model_residency(
            self.keep_face_models_on_device,
            total_memory_gib=total_memory_gib,
        )

    @staticmethod
    def _qwen_vram_config() -> dict[str, Any]:
        import torch

        return {
            "offload_dtype": "disk",
            "offload_device": "disk",
            "onload_dtype": torch.float8_e4m3fn,
            "onload_device": "cpu",
            "preparing_dtype": torch.float8_e4m3fn,
            "preparing_device": "cuda",
            "computation_dtype": torch.bfloat16,
            "computation_device": "cuda",
        }

    @staticmethod
    def _zimage_vram_config() -> dict[str, Any]:
        import torch

        return {
            "offload_dtype": torch.bfloat16,
            "offload_device": "cpu",
            "onload_dtype": torch.bfloat16,
            "onload_device": "cpu",
            "preparing_dtype": torch.bfloat16,
            "preparing_device": "cuda",
            "computation_dtype": torch.bfloat16,
            "computation_device": "cuda",
        }

    def _load_qwen(self) -> Any:
        if self._qwen_pipe is not None:
            return self._qwen_pipe
        self._require_cuda()
        os.environ.setdefault("DIFFSYNTH_DOWNLOAD_SOURCE", "huggingface")
        if self.hf_token:
            os.environ.setdefault("HF_TOKEN", self.hf_token)
        try:
            from diffsynth.pipelines.qwen_image import ControlNetInput, ModelConfig, QwenImagePipeline
        except ImportError as exc:
            raise ImportError(
                "The qwen-zimage pipeline needs the optional dependency group. "
                "Install: pip install 'remove-ai-watermarks[qwen-zimage]'"
            ) from exc

        self._progress("Loading Qwen-Image-2512, Lightning LoRA, and Canny ControlNet...")
        config = self._qwen_vram_config()
        model_configs = [
            ModelConfig(
                model_id=QWEN_IMAGE_2512_MODEL_ID,
                origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors",
                **config,
            ),
            ModelConfig(
                model_id=QWEN_IMAGE_2512_MODEL_ID,
                origin_file_pattern="text_encoder/model*.safetensors",
                **config,
            ),
            ModelConfig(
                model_id=QWEN_IMAGE_2512_MODEL_ID,
                origin_file_pattern="vae/diffusion_pytorch_model.safetensors",
                **config,
            ),
            ModelConfig(
                model_id=QWEN_CANNY_CONTROLNET_MODEL_ID,
                origin_file_pattern="model.safetensors",
                **config,
            ),
        ]
        pipe = QwenImagePipeline.from_pretrained(
            torch_dtype=self.torch_dtype,
            device=self.device,
            model_configs=model_configs,
            tokenizer_config=ModelConfig(
                model_id=QWEN_IMAGE_2512_MODEL_ID,
                origin_file_pattern="tokenizer/",
            ),
            vram_limit=self._vram_limit(),
        )
        from diffsynth.diffusion import FlowMatchScheduler

        # Avoid the base Qwen scheduler's terminal rescale for the distilled LoRA.
        # With exponential_shift_mu=log(3), this is the closest DiffSynth equivalent
        # of the source graph's four-step sgm_uniform + AuraFlow shift 3.
        pipe.scheduler = FlowMatchScheduler("Qwen-Image-Lightning")
        lightning = ModelConfig(
            model_id=QWEN_LIGHTNING_MODEL_ID,
            origin_file_pattern=QWEN_LIGHTNING_PATTERN,
        )
        pipe.load_lora(pipe.dit, lightning, alpha=0.8)
        if self.cache_prompt_embeddings:
            _cache_static_prompt_embeddings(
                pipe,
                ("prompt_emb", "prompt_emb_mask"),
            )
        self._qwen_pipe = (pipe, ControlNetInput)
        return self._qwen_pipe

    def _load_zimage(self) -> Any:
        if self._zimage_pipe is not None:
            return self._zimage_pipe
        self._require_cuda()
        os.environ.setdefault("DIFFSYNTH_DOWNLOAD_SOURCE", "huggingface")
        if self.hf_token:
            os.environ.setdefault("HF_TOKEN", self.hf_token)
        try:
            from diffsynth.pipelines.z_image import ModelConfig, ZImagePipeline
        except ImportError as exc:
            raise ImportError(
                "The qwen-zimage pipeline needs the optional dependency group. "
                "Install: pip install 'remove-ai-watermarks[qwen-zimage]'"
            ) from exc

        self._progress("Loading Z-Image Turbo face-detail model...")
        keep_on_device = self._keep_face_models_resident()
        config = self._zimage_vram_config()
        pipe = ZImagePipeline.from_pretrained(
            torch_dtype=self.torch_dtype,
            device=self.device,
            model_configs=[
                ModelConfig(
                    model_id=ZIMAGE_TURBO_MODEL_ID,
                    origin_file_pattern="transformer/*.safetensors",
                    **config,
                ),
                ModelConfig(
                    model_id=ZIMAGE_TURBO_MODEL_ID,
                    origin_file_pattern="text_encoder/*.safetensors",
                    **config,
                ),
                ModelConfig(
                    model_id=ZIMAGE_TURBO_MODEL_ID,
                    origin_file_pattern="vae/diffusion_pytorch_model.safetensors",
                    **config,
                ),
            ],
            tokenizer_config=ModelConfig(
                model_id=ZIMAGE_TURBO_MODEL_ID,
                origin_file_pattern="tokenizer/",
            ),
            vram_limit=self._vram_limit(),
        )
        if keep_on_device:
            _pin_vram_managed_models(pipe)
        if self.cache_prompt_embeddings:
            _cache_static_prompt_embeddings(pipe, ("prompt_embeds",))
        self._zimage_pipe = pipe
        return pipe

    def _load_sam(self) -> tuple[Any, Any]:
        if self._sam_model is not None and self._sam_processor is not None:
            return self._sam_model, self._sam_processor
        self._progress("Loading SAM face-mask model...")
        try:
            from transformers import AutoModelForMaskGeneration, AutoProcessor
        except ImportError as exc:
            raise ImportError("SAM needs transformers and torchvision from the qwen-zimage extra.") from exc
        kwargs: dict[str, Any] = {}
        if self.hf_token:
            kwargs["token"] = self.hf_token
        processor = AutoProcessor.from_pretrained(SAM_MODEL_ID, **kwargs)
        model = AutoModelForMaskGeneration.from_pretrained(
            SAM_MODEL_ID,
            torch_dtype=self.torch_dtype,
            **kwargs,
        ).to(self.device)
        model.eval()
        self._sam_model = model
        self._sam_processor = processor
        return model, processor

    def _sam_masks(
        self,
        image: Image.Image,
        boxes: list[tuple[int, int, int, int]],
    ) -> list[np.ndarray]:
        if not boxes:
            return []
        import torch

        try:
            model, processor = self._load_sam()
            input_points, input_labels = _sam_point_prompts(boxes)
            inputs = processor(
                images=image,
                input_boxes=[[list(box) for box in boxes]],
                input_points=input_points,
                input_labels=input_labels,
                return_tensors="pt",
            )
            original_sizes = inputs["original_sizes"].clone()
            reshaped_sizes = inputs["reshaped_input_sizes"].clone()
            inputs = _prepare_sam_inputs(inputs, self.device, self.torch_dtype)
            with torch.inference_mode():
                outputs = model(**inputs, multimask_output=True)
            processed = processor.post_process_masks(
                outputs.pred_masks.detach().cpu(),
                original_sizes,
                reshaped_sizes,
            )[0]
            mask_array, score_array = _sam_outputs_to_numpy(processed, outputs.iou_scores)
            binary_masks = _select_sam_masks(
                mask_array,
                score_array,
            )
            return _clip_sam_masks_to_boxes(binary_masks, boxes, image.size)
        except Exception as exc:
            log.warning("SAM face-mask refinement failed (%s); using box-derived ellipse masks", exc)
            return _ellipse_masks(boxes, image.size)

    def preload(self, *, global_only: bool = False) -> None:
        """Eagerly load the mandatory stage and, by default, the face stack."""
        self._load_qwen()
        _yunet_model_path()
        if not global_only:
            self._load_zimage()
            self._load_sam()

    def _run_global(self, image: Image.Image, strength: float, seed: int | None) -> Image.Image:
        pipe, controlnet_input_cls = self._load_qwen()
        input_image = _resize_to_target(image)
        control = build_canny_control_image(input_image)
        control_input = controlnet_input_cls(
            image=control,
            scale=float(self.controlnet_conditioning_scale),
        )
        self._progress(f"Running Qwen-Image-2512 Canny pass: strength={strength:.4f}, steps={GLOBAL_STEPS}...")
        result = pipe(
            **build_global_kwargs(
                input_image,
                strength=strength,
                seed=seed,
                controlnet_input=control_input,
            )
        )
        if result.size != image.size:
            result = result.resize(image.size, Image.Resampling.LANCZOS)
        return result.convert("RGB")

    @staticmethod
    def _detail_size(
        crop_size: tuple[int, int],
        face_size: tuple[int, int],
    ) -> tuple[int, int]:
        """Scale a crop toward a 768px face guide while capping it at 1024px."""
        crop_width, crop_height = crop_size
        face_width, face_height = face_size
        scale_for_face = 768.0 / max(1, max(face_width, face_height))
        scale_for_crop = 1024.0 / max(1, max(crop_width, crop_height))
        scale = min(scale_for_face, scale_for_crop)
        # Never shrink below the crop's current size unless the 1024 cap requires it.
        if max(crop_width, crop_height) <= 1024:
            scale = max(1.0, scale)
        width = max(16, round(crop_width * scale / 16.0) * 16)
        height = max(16, round(crop_height * scale / 16.0) * 16)
        return width, height

    def _run_faces(
        self,
        original: Image.Image,
        global_result: Image.Image,
        boxes: list[tuple[int, int, int, int]],
        masks: list[np.ndarray],
        *,
        strength: float,
        seed: int | None,
    ) -> Image.Image:
        if not boxes:
            return global_result
        pipe = self._load_zimage()
        base = np.asarray(global_result.convert("RGB")).copy()
        source = np.asarray(original.convert("RGB"))
        detail_seed = None if seed is None else seed + 1

        for index, (box, mask) in enumerate(zip(boxes, masks, strict=True), start=1):
            crop_box = _expanded_box(box, original.size)
            cx1, cy1, cx2, cy2 = crop_box
            crop_source = source[cy1:cy2, cx1:cx2]
            crop_mask = mask[cy1:cy2, cx1:cx2]
            if crop_source.size == 0 or not np.any(crop_mask):
                continue

            face_width = box[2] - box[0]
            face_height = box[3] - box[1]
            process_size = self._detail_size((cx2 - cx1, cy2 - cy1), (face_width, face_height))
            crop_image = Image.fromarray(crop_source).resize(process_size, Image.Resampling.LANCZOS)
            self._progress(
                f"Regenerating face {index}/{len(boxes)} with Z-Image: strength={strength:.4f}, steps={FACE_STEPS}..."
            )
            detailed = pipe(**build_face_kwargs(crop_image, strength=strength, seed=detail_seed))
            detailed = detailed.convert("RGB").resize((cx2 - cx1, cy2 - cy1), Image.Resampling.LANCZOS)

            base_crop = base[cy1:cy2, cx1:cx2]
            base[cy1:cy2, cx1:cx2] = composite_face(
                base_crop,
                np.asarray(detailed),
                crop_mask,
                feather=10,
            )
        return Image.fromarray(base)

    def run(
        self,
        image: Image.Image,
        *,
        strength: float | None,
        seed: int | None,
        tile: bool = False,
        tile_size: int = 1024,
        tile_overlap: int = 128,
    ) -> Image.Image:
        """Execute global regeneration and masked face repair."""
        self._require_cuda()
        seed = resolve_seed(seed, "qwen-zimage")
        global_strength = (
            resolution_adaptive_denoise(image.width, image.height) if strength is None else float(strength)
        )
        if tile and max(image.size) > tile_size:
            from remove_ai_watermarks.noai.tiling import run_tiled

            global_result = run_tiled(
                lambda tile_image: self._run_global(tile_image, global_strength, seed),
                image,
                tile_size,
                tile_overlap,
                self._progress,
            )
        else:
            global_result = self._run_global(image, global_strength, seed)

        self._progress("Detecting faces on the original image...")
        boxes = detect_faces(image)
        if not boxes:
            self._progress("No faces detected; keeping the Qwen global result.")
            return global_result
        masks = self._sam_masks(image, boxes)
        face_strength = largest_face_denoise(boxes, image.size) * FACE_DENOISE_SCALE
        return self._run_faces(
            image,
            global_result,
            boxes,
            masks,
            strength=face_strength,
            seed=seed,
        )
