import os
import cv2
import tempfile
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from pathlib import Path
import gradio as gr
import numpy as np

# Load the model into memory to make running multiple predictions efficient
img_colorization = pipeline(Tasks.image_colorization, model='iic/cv_ddcolor_image-colorization')

def inference(img):
    image = cv2.imread(str(img))
    output = img_colorization(image[..., ::-1])
    result = output[OutputKeys.OUTPUT_IMG].astype(np.uint8)

    temp_dir = tempfile.mkdtemp()
    out_path = os.path.join(temp_dir, 'old-to-color.png')
    cv2.imwrite(out_path, result)
    return Path(out_path)

# Modernized UI using Gradio 3.9 components
title = "🌈 Color Restorization Model"
description = "Upload a black & white photo to restore it in color using a deep learning model."

with gr.Blocks(title=title) as demo:
    gr.Markdown(f"## {title}")
    gr.Markdown(description)

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="filepath", label="Upload B&W Image")
            submit_btn = gr.Button("Colorize")
        with gr.Column():
            output_image = gr.Image(type="pil", label="Colorized Output")

    submit_btn.click(fn=inference, inputs=input_image, outputs=output_image)

demo.launch(enable_queue=True)