import gradio as gr
import torch
from transformers import AutoModelForImageToImage, AutoImageProcessor
from PIL import Image
import numpy as np
import os

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

# Load model components separately
image_processor = AutoImageProcessor.from_pretrained(
    "camenduru/cv_ddcolor_image-colorization",
    cache_dir=cache_dir
)
model = AutoModelForImageToImage.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 grayscale input
    if input_image.mode != 'L':
        input_image = input_image.convert('L')
    
    # Convert to RGB for model input
    rgb_image = input_image.convert("RGB")
    
    # Process through model
    with torch.inference_mode():
        # Preprocess
        pixel_values = image_processor(rgb_image, return_tensors="pt").pixel_values.to("cuda")
        
        # Forward pass
        outputs = model(pixel_values=pixel_values)
        
        # Postprocess
        output_image = image_processor.post_process(outputs, output_type="pil")[0]
    
    return output_image

# UI Layout
with gr.Blocks(theme="soft") 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")
    
    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()