Spaces:
Runtime error
Runtime error
| 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() |