import os

from PIL import Image
from model_handler import MODELS

upscaler = MODELS.upscaler
upscaler4SD = MODELS.upscaler4SD

def enhance_image(
    input_image: Image,
):
    
    h, w = input_image.size
    max_scale_size = 1024
    if h > max_scale_size:
        w = int(w * max_scale_size / h)
        h = max_scale_size
    if w > max_scale_size:
        h = int(h * max_scale_size / w)
        w = max_scale_size

    if h != input_image.size[1] or w != input_image.size[0]:
        input_image = input_image.resize((w, h))

    if os.environ.get("TILING", False):
        tileSizeStr = os.environ.get("TILE_SIZE", 1024)
        tileSize = int(tileSizeStr)
        enhanced_image = upscaler(input_image, tiling=True, tile_width=tileSize, tile_height=tileSize)
    else:
        enhanced_image = upscaler(input_image)

    return enhanced_image

def enhance_sd_image(
    input_image: Image,
):
    if os.environ.get("TILING", False):
        tileSizeStr = os.environ.get("TILE_SIZE", 1024)
        tileSize = int(tileSizeStr)
        enhanced_image = upscaler4SD(input_image, tiling=True, tile_width=tileSize, tile_height=tileSize)
    else:
        enhanced_image = upscaler4SD(input_image)

    return enhanced_image