engines/sdxl_tile_engine.py

"""
SDXL + ControlNet-Tile engine — the lighter generative fallback.

Runs when the Space is deployed on hardware where SUPIR is too heavy (e.g. a
16 GB T4). This is still a *generative* path: the image is bicubically enlarged
to the target size, then an SDXL img2img pass conditioned by a tile ControlNet
regenerates plausible fine detail at the new resolution (the community-standard
"ControlNet Tile upscale" recipe) — not an ESRGAN sharpener.

Models (both public):
  * base:      stabilityai/stable-diffusion-xl-base-1.0 (fp16 variant)
  * control:   xinsir/controlnet-tile-sdxl-1.0 (SDXL tile ControlNet)

VRAM: roughly 12–16 GB in fp16 for a 2048x2048 output.
"""

from __future__ import annotations

import logging
import threading

import torch
from PIL import Image

from . import UpscaleEngine, env_float, env_int, gpu_context

log = logging.getLogger("upscale.sdxl_tile")

BASE_REPO = "stabilityai/stable-diffusion-xl-base-1.0"
CONTROL_REPO = "xinsir/controlnet-tile-sdxl-1.0"

MAX_EDGE = env_int("SDXL_TILE_MAX_EDGE", 2048)          # cap output long edge
STRENGTH = env_float("SDXL_TILE_STRENGTH", 0.4)         # img2img denoise amount
CONTROL_SCALE = env_float("SDXL_TILE_CONTROL_SCALE", 0.6)
STEPS = env_int("SDXL_TILE_STEPS", 25)
GUIDANCE = env_float("SDXL_TILE_GUIDANCE", 5.0)
POSITIVE = (
    "highly detailed, sharp focus, natural textures, realistic, "
    "professional photography, 8k, intricate details"
)
NEGATIVE = (
    "blurry, out of focus, lowres, jpeg artifacts, oversmoothed, "
    "painting, cartoon, 3d render, watermark, text"
)


class SdxlTileEngine(UpscaleEngine):
    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._pipe = None
        self._device = "cuda" if torch.cuda.is_available() else "cpu"

    def name(self) -> str:
        return "sdxl_tile"

    def _load(self):
        with self._lock:
            if self._pipe is not None:
                return
            from diffusers import (
                AutoencoderKL,
                ControlNetModel,
                StableDiffusionXLControlNetPipeline,
            )

            log.info("Loading ControlNet (%s) ...", CONTROL_REPO)
            controlnet = ControlNetModel.from_pretrained(
                CONTROL_REPO, torch_dtype=torch.float16
            )
            log.info("Loading SDXL base (%s) ...", BASE_REPO)
            vae = AutoencoderKL.from_pretrained(
                BASE_REPO, subfolder="vae", torch_dtype=torch.float16
            )
            pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
                BASE_REPO,
                controlnet=controlnet,
                vae=vae,
                torch_dtype=torch.float16,
                variant="fp16",
                use_safetensors=True,
            )
            pipe = pipe.to(self._device)
            pipe.enable_vae_tiling()
            pipe.enable_vae_slicing()
            # enable_attention_slicing keeps peak VRAM low on 16 GB cards.
            try:
                pipe.enable_attention_slicing()
            except Exception:
                pass
            self._pipe = pipe
            log.info("SDXL + ControlNet-Tile ready on %s", self._device)

    def load(self) -> None:
        with gpu_context():
            self._load()

    def upscale(self, image: Image.Image, scale: int) -> Image.Image:
        with gpu_context():
            self._load()
            pipe = self._pipe
            assert pipe is not None

            rgb = image.convert("RGB")
            w, h = rgb.size
            w, h = w * scale, h * scale
            # Keep the long edge within the model's comfortable range.
            if max(w, h) > MAX_EDGE:
                k = MAX_EDGE / max(w, h)
                w, h = round(w * k), round(h * k)
            w, h = w - w % 8, h - h % 8  # SDXL latents want /8 dims

            init = rgb.resize((w, h), Image.LANCZOS)

            result = pipe(
                prompt=POSITIVE,
                negative_prompt=NEGATIVE,
                image=init,                    # both the init image and the tile-control input
                strength=STRENGTH,
                controlnet_conditioning_scale=CONTROL_SCALE,
                num_inference_steps=STEPS,
                guidance_scale=GUIDANCE,
                generator=torch.Generator(device=self._device).manual_seed(0),
            ).images[0]

            return result.convert("RGB")