src/UltimateSDUpscale/UltimateSDUpscale.py

"""Ultimate SD Upscale - tiled upscaling with seam fix."""
from src.AutoEncoders import VariationalAE
from src.sample import sampling
from src.UltimateSDUpscale import USDU_upscaler, image_util
import torch
from PIL import ImageFilter, ImageDraw, Image
from enum import Enum
import math

state = USDU_upscaler.state


class UnsupportedModel(Exception):
    pass


class StableDiffusionProcessing:
    """Container for SD processing parameters."""
    def __init__(self, init_img: Image.Image, model, positive, negative, vae, seed, steps, cfg,
                 sampler_name, scheduler, denoise, upscale_by, uniform_tile_mode, callback=None):
        self.init_images = [init_img]
        self.image_mask = None
        self.mask_blur = 0
        self.inpaint_full_res_padding = 0
        self.width, self.height = init_img.width, init_img.height
        self.model, self.positive, self.negative, self.vae = model, positive, negative, vae
        self.seed, self.steps, self.cfg = seed, steps, cfg
        self.sampler_name, self.scheduler, self.denoise = sampler_name, scheduler, denoise
        self.init_size = (init_img.width, init_img.height)
        self.upscale_by, self.uniform_tile_mode = upscale_by, uniform_tile_mode
        self.extra_generation_params = {}
        self.callback = callback


class Processed:
    """Container for processed images."""
    def __init__(self, p, images, seed, info):
        self.images, self.seed, self.info = images, seed, info

    def infotext(self, p, index):
        return None


def fix_seed(p):
    pass


def process_images(p, pipeline=False):
    """Process tiles using inpainting."""
    image_mask = p.image_mask.convert("L")
    init_image = p.init_images[0]

    crop_region = image_util.get_crop_region(image_mask, p.inpaint_full_res_padding)
    x1, y1, x2, y2 = crop_region
    crop_ratio = (x2 - x1) / (y2 - y1)
    p_ratio = p.width / p.height
    
    if crop_ratio > p_ratio:
        target_width, target_height = x2 - x1, round((x2 - x1) / p_ratio)
    else:
        target_width, target_height = round((y2 - y1) * p_ratio), y2 - y1
    
    crop_region, _ = image_util.expand_crop(crop_region, image_mask.width, image_mask.height, target_width, target_height)
    tile_size = (p.width, p.height)

    if p.mask_blur > 0:
        image_mask = image_mask.filter(ImageFilter.GaussianBlur(p.mask_blur))

    tiles = [img.crop(crop_region) for img in USDU_upscaler.batch]
    initial_tile_size = tiles[0].size
    tiles = [t.resize(tile_size, Image.Resampling.LANCZOS) if t.size != tile_size else t for t in tiles]

    positive_cropped = image_util.crop_cond(p.positive, crop_region, p.init_size, init_image.size, tile_size)
    negative_cropped = image_util.crop_cond(p.negative, crop_region, p.init_size, init_image.size, tile_size)

    batched_tiles = torch.cat([image_util.pil_to_tensor(t) for t in tiles], dim=0)
    (latent,) = VariationalAE.VAEEncode().encode(p.vae, batched_tiles)
    
    # Auto-detect Flux for disabling multi-scale and setting correct flags
    model_sampling_obj = p.model.get_model_object("model_sampling")
    from src.sample.sampling import ModelSamplingFlux, ModelSamplingFlux2
    is_flux = isinstance(model_sampling_obj, (ModelSamplingFlux, ModelSamplingFlux2))
    is_flux2 = isinstance(model_sampling_obj, ModelSamplingFlux2)

    # Pass crop offsets for positional embedding coherence (Critical for Flux/DiT)
    model_options = p.model.model_options.copy()
    transformer_options = model_options.get("transformer_options", {}).copy()
    transformer_options["top"] = y1
    transformer_options["left"] = x1
    model_options["transformer_options"] = transformer_options

    (samples,) = sampling.common_ksampler(p.model, p.seed, p.steps, p.cfg, p.sampler_name, p.scheduler,
                                          positive_cropped, negative_cropped, latent, denoise=p.denoise, 
                                          pipeline=pipeline, flux=is_flux, flux2=is_flux2,
                                          model_options=model_options, callback=p.callback)
    (decoded,) = VariationalAE.VAEDecode().decode(p.vae, samples)

    for i, tile_sampled in enumerate([image_util.tensor_to_pil(decoded, j) for j in range(len(decoded))]):
        init_image = USDU_upscaler.batch[i]
        if tile_sampled.size != initial_tile_size:
            tile_sampled = tile_sampled.resize(initial_tile_size, Image.Resampling.LANCZOS)

        image_tile_only = Image.new("RGBA", init_image.size)
        image_tile_only.paste(tile_sampled, crop_region[:2])
        temp = image_tile_only.copy()
        temp.putalpha(image_mask.resize(temp.size))
        image_tile_only.paste(temp, image_tile_only)
        result = init_image.convert("RGBA")
        result.alpha_composite(image_tile_only)
        USDU_upscaler.batch[i] = result.convert("RGB")

    return Processed(p, [USDU_upscaler.batch[0]], p.seed, None)


class USDUMode(Enum):
    LINEAR, CHESS, NONE = 0, 1, 2


class USDUSFMode(Enum):
    NONE, BAND_PASS, HALF_TILE, HALF_TILE_PLUS_INTERSECTIONS = 0, 1, 2, 3


class USDUpscaler:
    """Main upscaler class."""
    def __init__(self, p, image, upscaler_index, save_redraw, save_seams_fix, tile_width, tile_height):
        self.p, self.image = p, image
        self.scale_factor = math.ceil(max(p.width, p.height) / max(image.width, image.height))
        self.upscaler = USDU_upscaler.sd_upscalers[upscaler_index]
        self.redraw = USDURedraw()
        self.redraw.tile_width = tile_width or tile_height
        self.redraw.tile_height = tile_height or tile_width
        self.seams_fix = USDUSeamsFix()
        self.seams_fix.tile_width = self.redraw.tile_width
        self.seams_fix.tile_height = self.redraw.tile_height
        self.initial_info = None
        self.rows = math.ceil(self.p.height / self.redraw.tile_height)
        self.cols = math.ceil(self.p.width / self.redraw.tile_width)

    def get_factor(self, num):
        if num == 1: return 2
        for f in [4, 3, 2]:
            if num % f == 0: return f
        return 0

    def get_factors(self):
        scales, current = [], 1
        while current < self.scale_factor:
            f = self.get_factor(self.scale_factor // current)
            scales.append(f)
            current *= f
        self.scales = enumerate(scales)

    def upscale(self):
        print(f"Canva: {self.p.width}x{self.p.height}, Image: {self.image.width}x{self.image.height}, Scale: {self.scale_factor}")
        self.get_factors()
        for idx, val in self.scales:
            print(f"Upscaling iteration {idx + 1} with scale factor {val}")
            self.image = self.upscaler.scaler.upscale(self.image, val, self.upscaler.data_path)
        self.image = self.image.resize((self.p.width, self.p.height), resample=Image.LANCZOS)

    def setup_redraw(self, mode, padding, mask_blur):
        self.redraw.mode = USDUMode(mode)
        self.redraw.enabled = self.redraw.mode != USDUMode.NONE
        self.redraw.padding = padding
        self.p.mask_blur = mask_blur

    def setup_seams_fix(self, padding, denoise, mask_blur, width, mode):
        self.seams_fix.padding, self.seams_fix.denoise = padding, denoise
        self.seams_fix.mask_blur, self.seams_fix.width = mask_blur, width
        self.seams_fix.mode = USDUSFMode(mode)
        self.seams_fix.enabled = self.seams_fix.mode != USDUSFMode.NONE

    def calc_jobs_count(self):
        global state
        redraw = (self.rows * self.cols) if self.redraw.enabled else 0
        seams = self.rows * (self.cols - 1) + (self.rows - 1) * self.cols
        state.job_count = redraw + seams

    def print_info(self):
        print(f"Tile: {self.redraw.tile_width}x{self.redraw.tile_height}, Grid: {self.rows}x{self.cols}")

    def add_extra_info(self):
        self.p.extra_generation_params.update({
            "Ultimate SD upscale upscaler": self.upscaler.name,
            "Ultimate SD upscale tile_width": self.redraw.tile_width,
            "Ultimate SD upscale tile_height": self.redraw.tile_height,
        })

    def process(self, pipeline):
        USDU_upscaler.state.begin()
        self.calc_jobs_count()
        self.result_images = []
        if self.redraw.enabled:
            self.image = self.redraw.start(self.p, self.image, self.rows, self.cols, pipeline)
            self.initial_info = self.redraw.initial_info
        self.result_images.append(self.image)
        if self.seams_fix.enabled:
            self.image = self.seams_fix.start(self.p, self.image, self.rows, self.cols, pipeline)
            self.initial_info = self.seams_fix.initial_info
            self.result_images.append(self.image)
        USDU_upscaler.state.end()


class USDURedraw:
    """Tile redraw functionality."""
    def init_draw(self, p, width, height):
        p.inpaint_full_res = True
        p.inpaint_full_res_padding = self.padding
        p.width = math.ceil((self.tile_width + self.padding) / 64) * 64
        p.height = math.ceil((self.tile_height + self.padding) / 64) * 64
        mask = Image.new("L", (width, height), "black")
        return mask, ImageDraw.Draw(mask)

    def calc_rectangle(self, xi, yi):
        return xi * self.tile_width, yi * self.tile_height, (xi + 1) * self.tile_width, (yi + 1) * self.tile_height

    def linear_process(self, p, image, rows, cols, pipeline=False):
        global state
        mask, draw = self.init_draw(p, image.width, image.height)
        for yi in range(rows):
            for xi in range(cols):
                if state.interrupted: break
                draw.rectangle(self.calc_rectangle(xi, yi), fill="white")
                p.init_images, p.image_mask = [image], mask
                processed = process_images(p, pipeline)
                draw.rectangle(self.calc_rectangle(xi, yi), fill="black")
                if processed.images: image = processed.images[0]
        p.width, p.height = image.width, image.height
        self.initial_info = processed.infotext(p, 0)
        return image

    def start(self, p, image, rows, cols, pipeline=False):
        self.initial_info = None
        return self.linear_process(p, image, rows, cols, pipeline)


class USDUSeamsFix:
    """Seam fixing functionality."""
    def init_draw(self, p):
        self.initial_info = None
        p.width = math.ceil((self.tile_width + self.padding) / 64) * 64
        p.height = math.ceil((self.tile_height + self.padding) / 64) * 64

    def half_tile_process(self, p, image, rows, cols, pipeline=False):
        global state
        self.init_draw(p)
        processed = None
        gradient = Image.linear_gradient("L")
        
        row_gradient = Image.new("L", (self.tile_width, self.tile_height), "black")
        row_gradient.paste(gradient.resize((self.tile_width, self.tile_height // 2), Image.BICUBIC), (0, 0))
        row_gradient.paste(gradient.rotate(180).resize((self.tile_width, self.tile_height // 2), Image.BICUBIC), (0, self.tile_height // 2))
        
        col_gradient = Image.new("L", (self.tile_width, self.tile_height), "black")
        col_gradient.paste(gradient.rotate(90).resize((self.tile_width // 2, self.tile_height), Image.BICUBIC), (0, 0))
        col_gradient.paste(gradient.rotate(270).resize((self.tile_width // 2, self.tile_height), Image.BICUBIC), (self.tile_width // 2, 0))

        p.denoising_strength, p.mask_blur = self.denoise, self.mask_blur

        for yi in range(rows - 1):
            for xi in range(cols):
                p.width, p.height = self.tile_width, self.tile_height
                p.inpaint_full_res, p.inpaint_full_res_padding = True, self.padding
                mask = Image.new("L", (image.width, image.height), "black")
                mask.paste(row_gradient, (xi * self.tile_width, yi * self.tile_height + self.tile_height // 2))
                p.init_images, p.image_mask = [image], mask
                processed = process_images(p, pipeline)
                if processed.images: image = processed.images[0]

        for yi in range(rows):
            for xi in range(cols - 1):
                p.width, p.height = self.tile_width, self.tile_height
                p.inpaint_full_res, p.inpaint_full_res_padding = True, self.padding
                mask = Image.new("L", (image.width, image.height), "black")
                mask.paste(col_gradient, (xi * self.tile_width + self.tile_width // 2, yi * self.tile_height))
                p.init_images, p.image_mask = [image], mask
                processed = process_images(p, pipeline)
                if processed.images: image = processed.images[0]

        p.width, p.height = image.width, image.height
        if processed: self.initial_info = processed.infotext(p, 0)
        return image

    def start(self, p, image, rows, cols, pipeline=False):
        return self.half_tile_process(p, image, rows, cols, pipeline)


class Script(USDU_upscaler.Script):
    """Main script runner."""
    def run(self, p, _, tile_width, tile_height, mask_blur, padding, seams_fix_width, seams_fix_denoise,
            seams_fix_padding, upscaler_index, save_upscaled_image, redraw_mode, save_seams_fix_image,
            seams_fix_mask_blur, seams_fix_type, target_size_type, custom_width, custom_height, custom_scale, pipeline=False):
        fix_seed(p)
        USDU_upscaler.torch_gc()
        p.do_not_save_grid = p.do_not_save_samples = True
        p.inpaint_full_res = False
        p.inpainting_fill, p.n_iter, p.batch_size = 1, 1, 1

        init_img = image_util.flatten(p.init_images[0], USDU_upscaler.opts.img2img_background_color)
        p.width = math.ceil((init_img.width * custom_scale) / 64) * 64
        p.height = math.ceil((init_img.height * custom_scale) / 64) * 64

        upscaler = USDUpscaler(p, init_img, upscaler_index, save_upscaled_image, save_seams_fix_image, tile_width, tile_height)
        upscaler.upscale()
        upscaler.setup_redraw(redraw_mode, padding, mask_blur)
        upscaler.setup_seams_fix(seams_fix_padding, seams_fix_denoise, seams_fix_mask_blur, seams_fix_width, seams_fix_type)
        upscaler.print_info()
        upscaler.add_extra_info()
        upscaler.process(pipeline)
        return Processed(p, upscaler.result_images, p.seed, upscaler.initial_info or "")


# Monkey-patch overrides
_old_init = USDUpscaler.__init__
def _new_init(self, p, image, upscaler_index, save_redraw, save_seams_fix, tile_width, tile_height):
    # Determine downscale factor from model (8 for SD, 16 for Flux)
    downscale_factor = 8
    try:
        latent_format = p.model.get_model_object("latent_format")
        if hasattr(latent_format, "downscale_factor"):
            downscale_factor = latent_format.downscale_factor
    except Exception:
        pass
        
    p.width = math.ceil((image.width * p.upscale_by) / downscale_factor) * downscale_factor
    p.height = math.ceil((image.height * p.upscale_by) / downscale_factor) * downscale_factor
    _old_init(self, p, image, upscaler_index, save_redraw, save_seams_fix, tile_width, tile_height)
USDUpscaler.__init__ = _new_init

_old_redraw = USDURedraw.init_draw
def _new_redraw(self, p, width, height):
    mask, draw = _old_redraw(self, p, width, height)
    
    downscale_factor = 8
    try:
        latent_format = p.model.get_model_object("latent_format")
        if hasattr(latent_format, "downscale_factor"):
            downscale_factor = latent_format.downscale_factor
    except Exception:
        pass
        
    p.width = math.ceil((self.tile_width + self.padding) / downscale_factor) * downscale_factor
    p.height = math.ceil((self.tile_height + self.padding) / downscale_factor) * downscale_factor
    return mask, draw
USDURedraw.init_draw = _new_redraw

_old_seams = USDUSeamsFix.init_draw
def _new_seams(self, p):
    _old_seams(self, p)
    
    downscale_factor = 8
    try:
        latent_format = p.model.get_model_object("latent_format")
        if hasattr(latent_format, "downscale_factor"):
            downscale_factor = latent_format.downscale_factor
    except Exception:
        pass
        
    p.width = math.ceil((self.tile_width + self.padding) / downscale_factor) * downscale_factor
    p.height = math.ceil((self.tile_height + self.padding) / downscale_factor) * downscale_factor
USDUSeamsFix.init_draw = _new_seams

_old_upscale = USDUpscaler.upscale
def _new_upscale(self):
    _old_upscale(self)
    USDU_upscaler.batch = [self.image] + [img.resize((self.p.width, self.p.height), Image.LANCZOS) for img in USDU_upscaler.batch[1:]]
USDUpscaler.upscale = _new_upscale

MAX_RESOLUTION = 8192
MODES = {"Linear": USDUMode.LINEAR, "Chess": USDUMode.CHESS, "None": USDUMode.NONE}
SEAM_FIX_MODES = {"None": USDUSFMode.NONE, "Band Pass": USDUSFMode.BAND_PASS, "Half Tile": USDUSFMode.HALF_TILE, "Half Tile + Intersections": USDUSFMode.HALF_TILE_PLUS_INTERSECTIONS}


class UltimateSDUpscale:
    """Main entry point for Ultimate SD Upscale."""
    def upscale(self, image, model, positive, negative, vae, upscale_by, seed, steps, cfg, sampler_name,
                scheduler, denoise, upscale_model, mode_type, tile_width, tile_height, mask_blur, tile_padding,
                seam_fix_mode, seam_fix_denoise, seam_fix_mask_blur, seam_fix_width, seam_fix_padding,
                force_uniform_tiles, pipeline=False, callback=None):
        USDU_upscaler.sd_upscalers[0] = USDU_upscaler.UpscalerData()
        USDU_upscaler.actual_upscaler = upscale_model
        USDU_upscaler.batch = [image_util.tensor_to_pil(image, i) for i in range(len(image))]

        sdprocessing = StableDiffusionProcessing(
            image_util.tensor_to_pil(image), model, positive, negative, vae, seed, steps, cfg,
            sampler_name, scheduler, denoise, upscale_by, force_uniform_tiles, callback=callback)

        Script().run(
            p=sdprocessing, _=None, tile_width=tile_width, tile_height=tile_height, mask_blur=mask_blur,
            padding=tile_padding, seams_fix_width=seam_fix_width, seams_fix_denoise=seam_fix_denoise,
            seams_fix_padding=seam_fix_padding, upscaler_index=0, save_upscaled_image=False,
            redraw_mode=MODES[mode_type], save_seams_fix_image=False, seams_fix_mask_blur=seam_fix_mask_blur,
            seams_fix_type=SEAM_FIX_MODES[seam_fix_mode], target_size_type=2, custom_width=None,
            custom_height=None, custom_scale=upscale_by, pipeline=pipeline)

        return (torch.cat([image_util.pil_to_tensor(img) for img in USDU_upscaler.batch], dim=0),)