Update app.py

app.py

@@ -1,19 +1,21 @@
 import gradio as gr
 import torch
-from diffusers import DDColorPipeline
+from diffusers import AutoPipelineForImage2Image
+from transformers import pipeline
 from PIL import Image
 import numpy as np
 import os
 
-# Model loading with optimized settings
+# Model loading with dynamic pipeline selection
 cache_dir = "./model_cache"
 os.makedirs(cache_dir, exist_ok=True)
 
-# Load model once at startup
-pipe = DDColorPipeline.from_pretrained(
+# Load model using AutoPipeline
+pipe = AutoPipelineForImage2Image.from_pretrained(
     "camenduru/cv_ddcolor_image-colorization",
     torch_dtype=torch.float16,
-    cache_dir=cache_dir
+    cache_dir=cache_dir,
+    variant="fp16"
 ).to("cuda")
 
 def colorize_image(input_image):
@@ -22,16 +24,21 @@ def colorize_image(input_image):
     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)
+    # 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 numpy array for pipeline input
-    grayscale_array = np.array(resized_image)
+    # Convert to RGB as required by model
+    grayscale_image = resized_image.convert("RGB")
     
     # Generate colorized image
     with torch.inference_mode():
-        result = pipe(grayscale_array).images[0]
+        result = pipe(
+            prompt="colorized photo",
+            image=grayscale_image,
+            num_inference_steps=20,
+            strength=0.8
+        ).images[0]
     
     return result