Make CCO model default for colorization - Change default model from GAN (dummy) to CCO (working) - Add fallback to GAN if CCO not available - Fix image format and saving issues

app/main.py

@@ -122,6 +122,10 @@ def colorize_image_cco(img: Image.Image, model_name: str = "eccv16"):
     if model is None:
         raise ValueError(f"CCO model '{model_name}' not loaded")
     
+    # Ensure image is RGB
+    if img.mode != "RGB":
+        img = img.convert("RGB")
+    
     # Convert PIL Image to numpy array
     oimg = np.asarray(img)
     if oimg.ndim == 2:
@@ -134,11 +138,15 @@ def colorize_image_cco(img: Image.Image, model_name: str = "eccv16"):
     with torch.no_grad():
         out_ab = model(tens_l_rs)
     
-    # Postprocess output
+    # Postprocess output (returns RGB in [0, 1] range)
     output_rgb = postprocess_tens(tens_l_orig, out_ab)
     
+    # Clamp values to [0, 1] and convert to uint8
+    output_rgb = np.clip(output_rgb, 0, 1)
+    output_array = (output_rgb * 255).astype(np.uint8)
+    
     # Convert numpy array back to PIL Image
-    output_img = Image.fromarray((output_rgb * 255).astype(np.uint8))
+    output_img = Image.fromarray(output_array, 'RGB')
     return output_img
 
 def colorize_image(img: Image.Image, model_type: str = "gan", cco_model: str = "eccv16"):
@@ -298,7 +306,7 @@ async def colorize(
     user_id: Optional[str] = Form(None),
     category_id: Optional[str] = Form(None),
     categoryId: Optional[str] = Form(None),
-    model: Optional[str] = Form("gan"),  # New parameter: "gan", "cco", "cco-eccv16", "cco-siggraph17"
+    model: Optional[str] = Form(None),  # Model parameter: "gan", "cco", "cco-eccv16", "cco-siggraph17" (default: CCO if available)
 ):
     import time
     start_time = time.time()
@@ -314,13 +322,22 @@ async def colorize(
         effective_category_id = None
 
     # Parse model parameter
-    model_type = "gan"  # Default
-    cco_model = "eccv16"  # Default for CCO
-    model_type_for_log = "gan"  # For MongoDB logging
+    # Default to CCO if available, otherwise fallback to GAN
+    if CCO_AVAILABLE:
+        model_type = "cco"
+        cco_model = "eccv16"
+        model_type_for_log = "cco-eccv16"
+    else:
+        model_type = "gan"
+        model_type_for_log = "gan"
     
     if model:
         model = model.strip().lower()
-        if model == "cco" or model.startswith("cco-"):
+        if model == "gan":
+            # Use GAN model (dummy implementation - doesn't actually colorize)
+            model_type = "gan"
+            model_type_for_log = "gan"
+        elif model == "cco" or model.startswith("cco-"):
             if not CCO_AVAILABLE:
                 error_msg = "CCO models are not available"
                 log_api_call(
@@ -353,9 +370,8 @@ async def colorize(
                 cco_model = "eccv16"
                 model_type_for_log = "cco-eccv16"
         else:
-            # Default to "gan" for any other value
-            model_type = "gan"
-            model_type_for_log = "gan"
+            # Unknown model, use default (CCO if available, else GAN)
+            pass
 
     if not file.content_type.startswith("image/"):
         error_msg = "Invalid file type"
@@ -380,13 +396,21 @@ async def colorize(
 
     try:
         img = Image.open(io.BytesIO(await file.read()))
+        # Ensure image is RGB
+        if img.mode != "RGB":
+            img = img.convert("RGB")
+        
         output_img = colorize_image(img, model_type=model_type, cco_model=cco_model)
+        
+        # Ensure output is RGB
+        if output_img.mode != "RGB":
+            output_img = output_img.convert("RGB")
 
         processing_time = time.time() - start_time
 
         result_id = f"{uuid.uuid4()}.jpg"
         output_path = os.path.join(RESULTS_DIR, result_id)
-        output_img.save(output_path)
+        output_img.save(output_path, "JPEG", quality=95)
 
         base_url = "https://logicgoinfotechspaces-text-guided-image-colorization.hf.space"