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

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


def load_image_from_url(url):
    response = requests.get(url)
    img = Image.open(BytesIO(response.content))
    return img

def inference(img, img_url=None):
    if img_url:
        img = load_image_from_url(img_url)
    img = np.array(img)
    output = img_colorization(img[..., ::-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)
    upload_url = "https://api.postimages.org/upload"
    files = {'file': open(out_path, 'rb')}
    response = requests.post(upload_url, files=files)
    files.close()
    image_url = response.json()['url']  # رابط الصورة المحملة

    return Path(out_path), image_url

title = "Color Restorization Model"
interface = gr.Interface(
    inference,
    inputs=[
        gr.inputs.Image(type="pil", label="Input Image"),
        gr.inputs.Textbox(placeholder="Enter Image URL (optional)", label="Image URL (optional)")
    ],
    outputs=[
        gr.outputs.Image(type="pil", label="Output Image"),
        gr.outputs.Textbox(label="Download Link")
    ],
    title=title
)

interface.launch(enable_queue=True)