Update app.py

app.py

@@ -1,34 +1,72 @@
-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 torch
+from diffusers import DDColorPipeline
+from PIL import Image
 import numpy as np
+import os
 
-"""Load the model into memory to make running multiple predictions efficient"""
-img_colorization = pipeline(Tasks.image_colorization, model='iic/cv_ddcolor_image-colorization')
-
+# Model loading with optimized settings
+cache_dir = "./model_cache"
+os.makedirs(cache_dir, exist_ok=True)
 
-def inference(img):
-    image = cv2.imread(str(img))
+# Load model once at startup
+pipe = DDColorPipeline.from_pretrained(
+    "camenduru/cv_ddcolor_image-colorization",
+    torch_dtype=torch.float16,
+    cache_dir=cache_dir
+).to("cuda")
 
-    output = img_colorization(image[..., ::-1])
-    result = output[OutputKeys.OUTPUT_IMG].astype(np.uint8)
+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
 
-    temp_dir = tempfile.mkdtemp()
-    out_path = os.path.join(temp_dir, 'old-to-color.png')
-    cv2.imwrite(out_path, result)
-    return Path(out_path)
+# 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*
+    """)
 
-title = "Color Restorization Model"
-gr.Interface(
-    inference,
-    [gr.inputs.Image(type="filepath", label="Input")],
-    gr.outputs.Image(type="pil", label="Output"),
-    title=title
-).launch(enable_queue=True)
+demo.launch()