Update app.py

app.py

@@ -1,54 +1,49 @@
 import gradio as gr
 import torch
-from diffusers import AutoPipelineForImage2Image
-from transformers import pipeline
+from transformers import AutoModelForImageToImage, AutoImageProcessor
 from PIL import Image
 import numpy as np
 import os
 
-# Model loading with dynamic pipeline selection
+# Model loading with manual configuration
 cache_dir = "./model_cache"
 os.makedirs(cache_dir, exist_ok=True)
 
-# Load model using AutoPipeline
-pipe = AutoPipelineForImage2Image.from_pretrained(
+# 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,
-    variant="fp16"
+    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
+    # Ensure grayscale input
     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")
+    # Convert to RGB for model input
+    rgb_image = input_image.convert("RGB")
     
-    # Generate colorized image
+    # Process through model
     with torch.inference_mode():
-        result = pipe(
-            prompt="colorized photo",
-            image=grayscale_image,
-            num_inference_steps=20,
-            strength=0.8
-        ).images[0]
+        # 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 result
-
-# Custom CSS for vintage styling
-custom_css = """
-#output-image {max-width: 100%; border: 2px solid #ccc; border-radius: 8px;}
-"""
+    return output_image
 
 # UI Layout
-with gr.Blocks(theme="soft", css=custom_css) as demo:
+with gr.Blocks(theme="soft") as demo:
     gr.Markdown("""
     # 📸 Vintage Photo Colorizer
     Transform old black & white photos into vibrant color images using DDColor AI.
@@ -63,7 +58,7 @@ with gr.Blocks(theme="soft", css=custom_css) as demo:
         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")
+    output_img = gr.Image(label="Colorized Image")
     
     colorize_btn.click(
         fn=colorize_image,