import gradio as gr
import torch
from diffusers import DDColorPipeline
from PIL import Image
import numpy as np
import os

# Model loading with optimized settings
cache_dir = "./model_cache"
os.makedirs(cache_dir, exist_ok=True)

# Load model once at startup
pipe = DDColorPipeline.from_pretrained(
    "camenduru/cv_ddcolor_image-colorization",
    torch_dtype=torch.float16,
    cache_dir=cache_dir
).to("cuda")

def colorize_image(input_image):
    """Process B&W image and return colorized version"""
    # Ensure image is in grayscale mode
    if input_image.mode != 'L':
        input_image = input_image.convert('L')
    
    # Resize to model's expected input size (based on DDColor paper)
    target_size = (256, 256)
    resized_image = input_image.resize(target_size)
    
    # Convert to numpy array for pipeline input
    grayscale_array = np.array(resized_image)
    
    # Generate colorized image
    with torch.inference_mode():
        result = pipe(grayscale_array).images[0]
    
    return result

# Custom CSS for vintage styling
custom_css = """
#output-image {max-width: 100%; border: 2px solid #ccc; border-radius: 8px;}
"""

# UI Layout
with gr.Blocks(theme="soft", css=custom_css) as demo:
    gr.Markdown("""
    # 📸 Vintage Photo Colorizer
    Transform old black & white photos into vibrant color images using DDColor AI.
    
    ## How to Use
    1. Upload a grayscale image (or color image will be converted to B&W)
    2. Click "Colorize" to process
    3. Download your new colorized photo!
    """)
    
    with gr.Row():
        input_img = gr.Image(label="Upload Black & White Image", type="pil")
    
    colorize_btn = gr.Button("🎨 Colorize Photo", variant="primary")
    output_img = gr.Image(label="Colorized Image", elem_id="output-image")
    
    colorize_btn.click(
        fn=colorize_image,
        inputs=[input_img],
        outputs=[output_img]
    )
    
    gr.Markdown("""
    ### Powered by [DDColor](https://huggingface.co/papers/2212.11613 )  
    *Dual Decoders for Photo-Realistic Image Colorization*
    """)

demo.launch()