import gradio as gr
import torch
from diffusers import AutoPipelineForImage2Image
from transformers import pipeline
from PIL import Image
import numpy as np
import os

# Model loading with dynamic pipeline selection
cache_dir = "./model_cache"
os.makedirs(cache_dir, exist_ok=True)

# Load model using AutoPipeline
pipe = AutoPipelineForImage2Image.from_pretrained(
    "camenduru/cv_ddcolor_image-colorization",
    torch_dtype=torch.float16,
    cache_dir=cache_dir,
    variant="fp16"
).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
    target_size = (512, 512)  # Increased resolution for better quality
    resized_image = input_image.resize(target_size)
    
    # Convert to RGB as required by model
    grayscale_image = resized_image.convert("RGB")
    
    # Generate colorized image
    with torch.inference_mode():
        result = pipe(
            prompt="colorized photo",
            image=grayscale_image,
            num_inference_steps=20,
            strength=0.8
        ).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()