pipeline.py

"""VOSR inference helpers for the Gradio Space."""

from __future__ import annotations

import gc
import glob
import json
import logging
import os
import sys
import traceback
import types
from argparse import Namespace
from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.nn.functional as F
from huggingface_hub import snapshot_download
from PIL import Image
from safetensors.torch import load_file
from torchvision import transforms
from torchvision.transforms import Normalize

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [VOSR] %(levelname)s: %(message)s",
    force=True,
)
log = logging.getLogger("vosr")

ROOT = Path(__file__).resolve().parent
CKPT_ROOT = ROOT / "preset" / "ckpts"
HUB_MODEL = "CSWRY/VOSR"

torch.hub.set_dir(str(CKPT_ROOT / "torch_cache"))
sys.path.insert(0, str(ROOT))

from models.light_decoder import LightDecoder
from models.lightningdit import LightningDiT
from vosr import VOSR

IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)

SD2_LWDECODER_PATH = CKPT_ROOT / "sd21_lwdecoder.pth"
SD2_AE_PATH = CKPT_ROOT / "stable-diffusion-2-1-base"
QWEN_AE_PATH = CKPT_ROOT / "Qwen-Image-vae-2d"

MODEL_CHOICES = {
    "VOSR 1.4B multi-step (default)": {
        "id": "VOSR_1.4B_ms",
        "mode": "multistep",
        "default_steps": 25,
    },
    "VOSR 1.4B one-step": {
        "id": "VOSR_1.4B_os",
        "mode": "onestep",
        "default_steps": 1,
    },
    "VOSR 0.5B multi-step": {
        "id": "VOSR_0.5B_ms",
        "mode": "multistep",
        "default_steps": 25,
    },
    "VOSR 0.5B one-step": {
        "id": "VOSR_0.5B_os",
        "mode": "onestep",
        "default_steps": 1,
    },
}

DEFAULT_MODEL_LABEL = "VOSR 1.4B multi-step (default)"

_pipeline_cache: dict[str, dict[str, Any]] = {}
_assets_ready: set[str] = set()


def _clear_stale_distributed_env() -> None:
    for k in (
        "WORLD_SIZE",
        "RANK",
        "LOCAL_RANK",
        "LOCAL_WORLD_SIZE",
        "GROUP_RANK",
        "ROLE_RANK",
        "MASTER_ADDR",
        "MASTER_PORT",
        "TORCHELASTIC_RUN_ID",
        "TORCHELASTIC_MAX_RESTARTS",
    ):
        os.environ.pop(k, None)


def _accelerator_stub():
    _clear_stale_distributed_env()
    return types.SimpleNamespace(device=torch.device("cuda" if torch.cuda.is_available() else "cpu"))


def ensure_assets(model_id: str) -> Path:
    """Download only the assets needed for *model_id* into preset/ckpts."""
    if model_id in _assets_ready and (CKPT_ROOT / model_id / "args.json").is_file():
        return CKPT_ROOT / model_id

    CKPT_ROOT.mkdir(parents=True, exist_ok=True)
    args_probe = snapshot_download(
        HUB_MODEL,
        local_dir=str(CKPT_ROOT),
        allow_patterns=[f"{model_id}/args.json"],
    )
    with open(Path(args_probe) / model_id / "args.json") as f:
        meta = json.load(f)
    ae_type = meta.get("ae_type", "qwen")
    enc_type = meta.get("enc_type", "dinov2l")

    patterns = [
        f"{model_id}/**",
        "torch_cache/facebookresearch_dinov2_main/**",
    ]
    if enc_type == "dinov2l":
        patterns.append("torch_cache/checkpoints/dinov2_vitl14_pretrain.pth")
    else:
        patterns.append("torch_cache/checkpoints/dinov2_vitb14_pretrain.pth")

    if ae_type == "qwen":
        patterns.append("Qwen-Image-vae-2d/**")
    else:
        patterns.extend(
            [
                "sd21_lwdecoder.pth",
                "stable-diffusion-2-1-base/vae/config.json",
                "stable-diffusion-2-1-base/vae/diffusion_pytorch_model.safetensors",
            ]
        )

    print(f"Downloading assets for {model_id} from {HUB_MODEL}...")
    snapshot_download(HUB_MODEL, local_dir=str(CKPT_ROOT), allow_patterns=patterns)
    _assets_ready.add(model_id)
    return CKPT_ROOT / model_id


def load_config(ckpt_path: Path) -> Namespace:
    with open(ckpt_path / "args.json") as f:
        data = json.load(f)
    return Namespace(**data)


def preprocess_raw_image(x, args):
    x = x / 255.0
    x = F.interpolate(x, args.dinov2_size, mode="bicubic").clip(0.0, 1.0)
    return Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)(x)


def load_dinov2(args, device):
    if args.enc_type == "dinov2b":
        encoder = torch.hub.load("facebookresearch/dinov2", "dinov2_vitb14", trust_repo=True)
    elif args.enc_type == "dinov2l":
        encoder = torch.hub.load("facebookresearch/dinov2", "dinov2_vitl14", trust_repo=True)
    elif args.enc_type == "dinov2g":
        encoder = torch.hub.load("facebookresearch/dinov2", "dinov2_vitg14", trust_repo=True)
    else:
        raise ValueError(f"Unsupported enc_type: {args.enc_type}")

    del encoder.head
    encoder.head = torch.nn.Identity()

    def forward_with_features(self, x, masks=None):
        features = {}
        layer_indices = list(range(len(self.blocks)))
        if isinstance(x, list):
            return self.forward_features_list(x, masks)
        x = self.prepare_tokens_with_masks(x, masks)
        for i, blk in enumerate(self.blocks):
            x = blk(x)
            if i in layer_indices:
                features[f"layer_{i}"] = x[:, 1:]
        x_norm = self.norm(x)
        return features, x_norm[:, 1:]

    encoder.forward_with_features = types.MethodType(forward_with_features, encoder)
    encoder = encoder.to(device).eval()
    for p in encoder.parameters():
        p.requires_grad_(False)
    return encoder


def get_venc_features(venc, lq_tensor, args):
    with torch.no_grad():
        raw_image = (0.5 * lq_tensor + 0.5) * 255
        raw_image_ = preprocess_raw_image(raw_image, args)
        features, x_norm = venc.forward_with_features(raw_image_)
        z = [v for k, v in features.items() if k.startswith("layer_")]
        z[-1] = x_norm
        z = [z[i] for i in args.layer_dinov2b_list]
    return z


AE_FACTOR = 8


def _vram(tag: str) -> None:
    if not torch.cuda.is_available():
        return
    alloc = torch.cuda.memory_allocated() / 1e9
    reserved = torch.cuda.memory_reserved() / 1e9
    peak = torch.cuda.max_memory_allocated() / 1e9
    log.info("%s | VRAM alloc=%.2fGB reserved=%.2fGB peak=%.2fGB", tag, alloc, reserved, peak)


def _vae_blend_mask(tile_h: int, tile_w: int, channels: int, device, overlap_h: int = 0, overlap_w: int = 0):
    """
    Soft rectangular blend for VAE tiles (Hann / raised-cosine).

    The old peaked Gaussian (var=0.01) looked like circular masks with grey
    corners because edge weights were ~0 and neighbors did not cover them.
    """
    def _ramp(n: int, ov: int) -> torch.Tensor:
        if n <= 1:
            return torch.ones(n, device=device, dtype=torch.float32)
        # Flat center, cosine ramps only in the overlap margins.
        ov = int(max(0, min(ov, n // 2)))
        w = torch.ones(n, device=device, dtype=torch.float32)
        if ov <= 0:
            return w
        t = torch.linspace(0, torch.pi, ov + 2, device=device)[1:-1]
        ramp_up = 0.5 - 0.5 * torch.cos(t)  # 0 -> 1
        ramp_down = ramp_up.flip(0)  # 1 -> 0
        w[:ov] = ramp_up
        w[-ov:] = ramp_down
        return w

    wy = _ramp(tile_h, overlap_h)
    wx = _ramp(tile_w, overlap_w)
    w = wy[:, None] * wx[None, :]
    # Keep a small floor so uncovered edge pixels never explode / go grey.
    w = w.clamp_min(1e-2)
    return w.unsqueeze(0).unsqueeze(0).expand(1, channels, -1, -1)


def _tile_starts(length: int, tile: int, overlap: int) -> list[int]:
    stride = max(tile - overlap, 1)
    if length <= tile:
        return [0]
    positions = list(range(0, length - tile + 1, stride))
    if positions[-1] + tile < length:
        positions.append(length - tile)
    return sorted(set(positions))


def _vae_stats(vae, args, device):
    if args.ae_type == "qwen":
        latents_mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1).to(device)
        latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, -1, 1, 1).to(device)
        return latents_mean, latents_std
    return None, None


def _encode_once(vae, x, args, use_mode: bool):
    posterior = vae.encode(x).latent_dist
    z = posterior.mode() if use_mode else posterior.sample()
    if args.ae_type == "qwen":
        return z
    if args.ae_type == "sd2":
        return z * vae.config.scaling_factor
    raise ValueError(args.ae_type)


def _decode_once(vae, z, args, latents_mean, latents_std, light_decoder=None):
    if args.ae_type == "sd2":
        return light_decoder(z / vae.config.scaling_factor).clamp(-1, 1)
    if args.ae_type == "qwen":
        z = z / latents_std + latents_mean
        return vae.decode(z, return_dict=False)[0].clamp(-1, 1)
    raise ValueError(args.ae_type)


def _tiled_encode(vae, x, args, device):
    """Overlap-blend VAE encode in pixel space. Uses mode() for seam stability."""
    tile = int(getattr(args, "vae_tile_size", 0) or 0)
    overlap = int(getattr(args, "vae_tile_overlap", 64) or 0)
    _, _, h, w = x.shape
    if tile <= 0 or (h <= tile and w <= tile):
        log.info("VAE encode full | input=%sx%s", h, w)
        _vram("before VAE encode")
        latents_mean, latents_std = _vae_stats(vae, args, device)
        z = _encode_once(vae, x, args, use_mode=False)
        if args.ae_type == "qwen":
            z = (z - latents_mean) * latents_std
        _vram("after VAE encode")
        return z, latents_mean, latents_std

    tile = min(tile, h, w)
    # Need meaningful overlap; tiny values recreate seam / mask artifacts.
    min_ov = max(tile // 8, AE_FACTOR * 4)
    if overlap < min_ov:
        log.warning("VAE tile overlap %s too small for tile %s — raising to %s", overlap, tile, min_ov)
        overlap = min_ov
    overlap = min(overlap, tile // 2)
    # align tile to AE_FACTOR so latent tiles are integer
    tile = max((tile // AE_FACTOR) * AE_FACTOR, AE_FACTOR)
    overlap = max((overlap // AE_FACTOR) * AE_FACTOR, 0)

    latents_mean, latents_std = _vae_stats(vae, args, device)
    h_pos = _tile_starts(h, tile, overlap)
    w_pos = _tile_starts(w, tile, overlap)
    log.info(
        "VAE encode tiled | input=%sx%s tile=%s overlap=%s grid=%sx%s",
        h, w, tile, overlap, len(h_pos), len(w_pos),
    )
    _vram("before VAE tiled encode")

    # probe channels
    probe = _encode_once(vae, x[:, :, :tile, :tile], args, use_mode=True)
    if args.ae_type == "qwen":
        probe = (probe - latents_mean) * latents_std
    lc = probe.shape[1]
    lh, lw = h // AE_FACTOR, w // AE_FACTOR
    lt = tile // AE_FACTOR
    lo = overlap // AE_FACTOR
    out = torch.zeros(1, lc, lh, lw, device=device, dtype=probe.dtype)
    weight = torch.zeros_like(out)
    g = _vae_blend_mask(lt, lt, lc, device, overlap_h=lo, overlap_w=lo)

    for hi in h_pos:
        for wi in w_pos:
            crop = x[:, :, hi : hi + tile, wi : wi + tile]
            z_tile = _encode_once(vae, crop, args, use_mode=True)
            if args.ae_type == "qwen":
                z_tile = (z_tile - latents_mean) * latents_std
            lhi, lwi = hi // AE_FACTOR, wi // AE_FACTOR
            out[:, :, lhi : lhi + lt, lwi : lwi + lt] += z_tile * g
            weight[:, :, lhi : lhi + lt, lwi : lwi + lt] += g

    z = out / weight.clamp_min(1e-6)
    _vram("after VAE tiled encode")
    log.info("VAE encode done | latent=%sx%s", z.shape[-2], z.shape[-1])
    return z, latents_mean, latents_std


def _tiled_decode(vae, sr_latent, args, latents_mean, latents_std, light_decoder=None):
    tile_px = int(getattr(args, "vae_tile_size", 0) or 0)
    overlap_px = int(getattr(args, "vae_tile_overlap", 64) or 0)
    _, lc, lh, lw = sr_latent.shape

    if tile_px <= 0:
        log.info("VAE decode full | latent=%sx%s", lh, lw)
        _vram("before VAE decode")
        out = _decode_once(vae, sr_latent, args, latents_mean, latents_std, light_decoder)
        _vram("after VAE decode")
        return out

    lt = max(tile_px // AE_FACTOR, 1)
    lo = max(overlap_px // AE_FACTOR, 0)
    if lh <= lt and lw <= lt:
        log.info("VAE decode full (fits tile) | latent=%sx%s", lh, lw)
        return _decode_once(vae, sr_latent, args, latents_mean, latents_std, light_decoder)

    lt = min(lt, lh, lw)
    min_lo = max(lt // 8, 4)
    if lo < min_lo:
        log.warning("VAE decode overlap %s too small for lt=%s — raising to %s", lo, lt, min_lo)
        lo = min_lo
    lo = min(lo, lt // 2)
    h_pos = _tile_starts(lh, lt, lo)
    w_pos = _tile_starts(lw, lt, lo)
    log.info(
        "VAE decode tiled | latent=%sx%s lt=%s lo=%s grid=%sx%s",
        lh, lw, lt, lo, len(h_pos), len(w_pos),
    )
    _vram("before VAE tiled decode")

    probe = _decode_once(
        vae,
        sr_latent[:, :, :lt, :lt],
        args,
        latents_mean,
        latents_std,
        light_decoder,
    )
    oc, oh, ow = probe.shape[1], lh * AE_FACTOR, lw * AE_FACTOR
    pt = lt * AE_FACTOR
    po = lo * AE_FACTOR
    out = torch.zeros(1, oc, oh, ow, device=sr_latent.device, dtype=probe.dtype)
    weight = torch.zeros_like(out)
    g = _vae_blend_mask(pt, pt, oc, sr_latent.device, overlap_h=po, overlap_w=po)

    for hi in h_pos:
        for wi in w_pos:
            z_tile = sr_latent[:, :, hi : hi + lt, wi : wi + lt]
            rgb = _decode_once(vae, z_tile, args, latents_mean, latents_std, light_decoder)
            phi, pwi = hi * AE_FACTOR, wi * AE_FACTOR
            out[:, :, phi : phi + pt, pwi : pwi + pt] += rgb * g
            weight[:, :, phi : phi + pt, pwi : pwi + pt] += g

    rgb = out / weight.clamp_min(1e-6)
    _vram("after VAE tiled decode")
    return rgb


def _encode_latent(vae, x, args, device):
    return _tiled_encode(vae, x, args, device)


def _decode_latent(vae, sr_latent, args, latents_mean, latents_std, light_decoder=None):
    return _tiled_decode(vae, sr_latent, args, latents_mean, latents_std, light_decoder)


def adain_color_fix(target, source):
    from torchvision.transforms.functional import to_pil_image, to_tensor

    target_tensor = to_tensor(target).unsqueeze(0)
    source_tensor = to_tensor(source).unsqueeze(0)
    eps = 1e-5
    target_mean = torch.mean(target_tensor, dim=[2, 3], keepdim=True)
    target_std = torch.std(target_tensor, dim=[2, 3], keepdim=True) + eps
    source_mean = torch.mean(source_tensor, dim=[2, 3], keepdim=True)
    source_std = torch.std(source_tensor, dim=[2, 3], keepdim=True) + eps
    target_tensor = (target_tensor - target_mean) / target_std * source_std + source_mean
    return to_pil_image(torch.clamp(target_tensor[0], 0, 1))


def wavelet_color_fix(target, source):
    import cv2

    target_np = np.array(target).astype(np.float32) / 255.0
    source_np = np.array(source.resize(target.size, Image.LANCZOS)).astype(np.float32) / 255.0
    source_low = cv2.GaussianBlur(source_np, (0, 0), 5)
    target_low = cv2.GaussianBlur(target_np, (0, 0), 5)
    result = np.clip(source_low + (target_np - target_low), 0, 1) * 255.0
    return Image.fromarray(result.astype(np.uint8))


def _find_weight(ckpt_path: Path) -> Path:
    search_dirs = [ckpt_path / "clean_weights", ckpt_path / "checkpoints", ckpt_path]
    for d in search_dirs:
        for name in ("ema_model.safetensors", "model.safetensors"):
            candidate = d / name
            if candidate.is_file():
                return candidate
    hits = glob.glob(str(ckpt_path / "**" / "*.safetensors"), recursive=True)
    if not hits:
        raise FileNotFoundError(f"No .safetensors under {ckpt_path}")
    return Path(hits[0])


def load_pipeline(model_label: str) -> dict[str, Any]:
    if model_label in _pipeline_cache:
        return _pipeline_cache[model_label]

    spec = MODEL_CHOICES[model_label]
    ckpt_path = ensure_assets(spec["id"])
    args = load_config(ckpt_path)
    args.ae_path = str(QWEN_AE_PATH if args.ae_type == "qwen" else SD2_AE_PATH)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    accelerator = _accelerator_stub()
    print(f"Loading {model_label} on {device}...")

    if args.ae_type == "qwen":
        from models.qwenimage_vae2d import AutoencoderKLQwenImage2D

        vae = AutoencoderKLQwenImage2D.from_pretrained(args.ae_path)
    else:
        from diffusers import AutoencoderKL

        vae = AutoencoderKL.from_pretrained(args.ae_path, subfolder="vae")
    vae.to(device).eval()

    light_decoder = None
    if args.ae_type == "sd2":
        ckpt = torch.load(SD2_LWDECODER_PATH, map_location="cpu")
        dec_config = ckpt["config"]
        light_decoder = LightDecoder(
            in_channels=dec_config["in_channels"],
            out_channels=dec_config["out_channels"],
            block_out_channels=tuple(dec_config["block_out_channels"]),
            layers_per_block=dec_config["layers_per_block"],
        )
        light_decoder.load_state_dict(ckpt["model_state_dict"])
        light_decoder.to(device).eval()

    venc = load_dinov2(args, device)

    base_channel = 4 if args.ae_type == "sd2" else 16
    onestep = spec["mode"] == "onestep"
    model = LightningDiT(
        input_size=args.resolution // 8,
        patch_size=args.patch_size,
        in_channels=2 * base_channel,
        out_channels=base_channel,
        hidden_size=args.dim,
        depth=args.depth,
        num_heads=args.num_heads,
        mlp_ratio=args.mlp_ratio,
        z_dims=args.enc_dim,
        encdim_ratio=args.encdim_ratio,
        auxiliary_time_cond=onestep,
        use_qknorm=args.use_qknorm,
        use_swiglu=args.use_swiglu,
        use_rope=args.use_rope,
        use_rmsnorm=args.use_rmsnorm,
        wo_shift=args.wo_shift,
        num_fused_layers=len(args.layer_dinov2b_list),
    )
    weight_path = _find_weight(ckpt_path)
    print(f"Loading weights from {weight_path}")
    state_dict = load_file(str(weight_path))
    model.load_state_dict(state_dict, strict=False)
    model.to(device).eval()
    model.forward = model.forward_flexible

    vosr_model = VOSR(
        time_dist=args.time_dist,
        cfg_ratio=args.cfg_ratio,
        cfg_scale=getattr(args, "cfg_scale", 2.0),
        interp_type=args.interp_type,
        accelerator=accelerator,
        t_start=getattr(args, "t_start", None) or 0.0,
        t_end=getattr(args, "t_end", None) or 1.0,
        args=args,
    )

    pipe = {
        "label": model_label,
        "spec": spec,
        "args": args,
        "device": device,
        "vae": vae,
        "venc": venc,
        "model": model,
        "vosr_model": vosr_model,
        "light_decoder": light_decoder,
        "mode": spec["mode"],
    }
    _pipeline_cache[model_label] = pipe
    return pipe


def unload_other_pipelines(keep_label: str) -> None:
    """Free VRAM from pipelines that are not currently selected."""
    drop = [k for k in _pipeline_cache if k != keep_label]
    for k in drop:
        pipe = _pipeline_cache.pop(k)
        for key in ("vae", "venc", "model", "light_decoder", "vosr_model"):
            obj = pipe.get(key)
            if obj is None:
                continue
            try:
                obj.to("cpu")
            except Exception:
                pass
        del pipe
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def run_sr(
    image: Image.Image,
    model_label: str,
    upscale: int,
    infer_steps: int,
    cfg_scale: float,
    weak_cond: float,
    tile_size: int,
    tile_overlap: int,
    align_method: str,
    seed: int,
    vae_tile_size: int = 0,
    vae_tile_overlap: int = 64,
) -> Image.Image:
    if image is None:
        raise ValueError("Please upload an image.")

    if torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats()

    unload_other_pipelines(model_label)
    pipe = load_pipeline(model_label)
    args = pipe["args"]
    device = pipe["device"]
    model = pipe["model"]
    vae = pipe["vae"]
    venc = pipe["venc"]
    vosr_model = pipe["vosr_model"]
    light_decoder = pipe["light_decoder"]
    mode = pipe["mode"]

    args.infer_steps = int(infer_steps)
    args.tile_size = int(tile_size)
    args.tile_overlap = int(tile_overlap)
    args.vae_tile_size = int(vae_tile_size or 0)
    args.vae_tile_overlap = int(vae_tile_overlap or 0)
    args.upscale = int(upscale)
    args.weak_cond_strength_aelq_list = [float(weak_cond), float(weak_cond)]
    vosr_model.cfg_scale = float(cfg_scale)
    vosr_model.args = args

    torch.manual_seed(int(seed))
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(int(seed))
    np.random.seed(int(seed))

    raw_img = image.convert("RGB")
    w, h = raw_img.size
    target_w, target_h = w * int(upscale), h * int(upscale)
    # VAE stride requires multiples of 8
    pad_w = (AE_FACTOR - target_w % AE_FACTOR) % AE_FACTOR
    pad_h = (AE_FACTOR - target_h % AE_FACTOR) % AE_FACTOR
    work_w, work_h = target_w + pad_w, target_h + pad_h

    log.info(
        "run_sr start | model=%s mode=%s ae=%s | input=%sx%s upscale=%s -> target=%sx%s work=%sx%s "
        "(pad=%s,%s) | dit_tile=%s/%s vae_tile=%s/%s steps=%s cfg=%s",
        model_label,
        mode,
        getattr(args, "ae_type", "?"),
        w,
        h,
        upscale,
        target_w,
        target_h,
        work_w,
        work_h,
        pad_w,
        pad_h,
        args.tile_size,
        args.tile_overlap,
        args.vae_tile_size,
        args.vae_tile_overlap,
        args.infer_steps,
        cfg_scale,
    )
    if target_w != target_h and int(tile_size or 0) <= 0:
        log.warning(
            "Non-square target (%sx%s) without DiT tiling — forward_flexible requires H==W. "
            "Enable DiT tile size (e.g. 512).",
            target_w,
            target_h,
        )

    input_img = raw_img.resize((target_w, target_h), Image.BICUBIC)
    if pad_w or pad_h:
        padded = Image.new("RGB", (work_w, work_h))
        padded.paste(input_img, (0, 0))
        work_img = padded
    else:
        work_img = input_img
    lq = transforms.ToTensor()(work_img).unsqueeze(0).to(device) * 2.0 - 1.0
    log.info("LQ tensor shape=%s dtype=%s device=%s", tuple(lq.shape), lq.dtype, lq.device)
    _vram("after LQ to GPU")

    try:
        with torch.no_grad():
            if tile_size and int(tile_size) > 0:
                from inference_tiles import tiled_multistep, tiled_onestep

                log.info("DiT path: tiled")
                if mode == "multistep":
                    sr_tensor = tiled_multistep(
                        model, vosr_model, vae, venc, lq, args, device=device, light_decoder=light_decoder
                    )
                else:
                    sr_tensor = tiled_onestep(
                        model, vae, venc, lq, args, device=device, light_decoder=light_decoder
                    )
            else:
                log.info("DiT path: full-frame")
                lq_latent, latents_mean, latents_std = _encode_latent(vae, lq, args, device)
                log.info("latent shape=%s", tuple(lq_latent.shape))
                if lq_latent.shape[-2] != lq_latent.shape[-1]:
                    raise ValueError(
                        f"Latent is non-square {lq_latent.shape[-2]}x{lq_latent.shape[-1]}. "
                        "Enable DiT tiling (tile size > 0) for non-square images."
                    )
                z_fea = get_venc_features(venc, lq, args)
                _vram("after DINOv2 features")
                if mode == "multistep":
                    sr_latent = vosr_model.sample_multistep_fm(
                        model, lq_latent, n_steps=args.infer_steps, venc_fea=z_fea
                    )
                else:
                    sr_latent = vosr_model.sample_onestep(
                        model, lq_latent, n_steps=args.infer_steps, venc_fea=z_fea
                    )
                _vram("after DiT sample")
                sr_tensor = _decode_latent(vae, sr_latent, args, latents_mean, latents_std, light_decoder)
    except Exception as exc:
        log.error(
            "run_sr FAILED | model=%s target=%sx%s dit_tile=%s vae_tile=%s | %s: %s",
            model_label,
            target_w,
            target_h,
            args.tile_size,
            args.vae_tile_size,
            type(exc).__name__,
            exc,
        )
        log.error("traceback:\n%s", traceback.format_exc())
        raise RuntimeError(
            f"{type(exc).__name__}: {exc} "
            f"(input {w}x{h} → {target_w}x{target_h}, "
            f"dit_tile={args.tile_size}, vae_tile={args.vae_tile_size})"
        ) from exc

    # crop padding
    sr_img = transforms.ToPILImage()(sr_tensor[0].float().cpu() * 0.5 + 0.5)
    if pad_w or pad_h:
        sr_img = sr_img.crop((0, 0, target_w, target_h))
    if align_method == "adain":
        sr_img = adain_color_fix(sr_img, input_img)
    elif align_method == "wavelet":
        sr_img = wavelet_color_fix(sr_img, input_img)
    _vram("run_sr done")
    log.info("run_sr success | output=%sx%s", sr_img.size[0], sr_img.size[1])
    return sr_img