"""Gradio and ZeroGPU entry point. The model revision is pinned at release."""
import os
from pathlib import Path
import spaces
import torch
import gradio as gr
from huggingface_hub import snapshot_download
from inference import load_colorizer, prepare, chroma_coefficients, render

REPO=os.getenv('MODEL_ID','User-2468/mini-unet-colorizer')
REV=os.getenv('MODEL_REVISION','v3.0.0')
CPU=os.getenv('ZEROGPU_CPU_TEST')=='1'
DEVICE='cpu' if CPU else 'cuda'
if CPU:torch.set_num_threads(min(4,os.cpu_count() or 1))
path=Path(REPO)
if not path.is_dir():path=Path(snapshot_download(REPO,revision=REV,allow_patterns=['model.safetensors','config.json']))
MODEL=load_colorizer(path,DEVICE)

@spaces.GPU(duration=20)
def infer(small,radius):
    return chroma_coefficients(MODEL,small,int(radius))

def run(image,strength,smoothing):
    if image is None:raise gr.Error('Upload a photograph first.')
    radius={'Balanced':8,'Gentle':4,'Strong':16}[smoothing]
    try:
        light,alpha,small=prepare(image,256)
        coefficients=infer(small,radius)
        result=render(light,alpha,coefficients,strength)
    except (ValueError,RuntimeError,TypeError) as exc:
        raise gr.Error(str(exc)) from exc
    return result

with gr.Blocks(title='Mini Photo Colorizer') as demo:
    gr.Markdown('# Mini Photo Colorizer\nAdd plausible colour to a black-and-white photograph. Original colours may be unknowable.')
    with gr.Row():
        source=gr.Image(type='pil',image_mode='RGBA',label='Original photograph',sources=['upload','clipboard'],height=440)
        result=gr.Image(type='pil',label='Colourised photograph',format='png',height=440)
    with gr.Row():
        strength=gr.Slider(0,1.5,value=1.,step=.05,label='Colour strength')
        smoothing=gr.Radio(['Gentle','Balanced','Strong'],value='Balanced',label='Colour smoothing')
    button=gr.Button('Colourise',variant='primary')
    button.click(run,[source,strength,smoothing],result,api_name='colorize',concurrency_limit=1)
    gr.ClearButton([source,result])
    gr.Markdown('3.995 million parameters · Full-resolution output, up to 12 megapixels · PNG download\n\nBest suited to ordinary photographs. Colours can be muted or inaccurate on unusual scenes, tiny objects and heavily damaged scans.')
demo.queue(default_concurrency_limit=1,max_size=12)
if __name__=='__main__':demo.launch()