Integrate CCO colorization models (eccv16 and siggraph17) - Add CCO colorizers module from kinsung/cco - Update /colorize endpoint to support model selection parameter - Add scikit-image dependency - Maintain backward compatibility with existing GAN model - Update MongoDB logging to track model type used

app/colorizers/__init__.py

@@ -0,0 +1,6 @@
+
+from .base_color import *
+from .eccv16 import *
+from .siggraph17 import *
+from .util import *
+

app/colorizers/base_color.py

@@ -0,0 +1,24 @@
+
+import torch
+from torch import nn
+
+class BaseColor(nn.Module):
+	def __init__(self):
+		super(BaseColor, self).__init__()
+
+		self.l_cent = 50.
+		self.l_norm = 100.
+		self.ab_norm = 110.
+
+	def normalize_l(self, in_l):
+		return (in_l-self.l_cent)/self.l_norm
+
+	def unnormalize_l(self, in_l):
+		return in_l*self.l_norm + self.l_cent
+
+	def normalize_ab(self, in_ab):
+		return in_ab/self.ab_norm
+
+	def unnormalize_ab(self, in_ab):
+		return in_ab*self.ab_norm
+

app/colorizers/eccv16.py

@@ -0,0 +1,105 @@
+
+import torch
+import torch.nn as nn
+import numpy as np
+from IPython import embed
+
+from .base_color import *
+
+class ECCVGenerator(BaseColor):
+    def __init__(self, norm_layer=nn.BatchNorm2d):
+        super(ECCVGenerator, self).__init__()
+
+        model1=[nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1, bias=True),]
+        model1+=[nn.ReLU(True),]
+        model1+=[nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=True),]
+        model1+=[nn.ReLU(True),]
+        model1+=[norm_layer(64),]
+
+        model2=[nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=True),]
+        model2+=[nn.ReLU(True),]
+        model2+=[nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1, bias=True),]
+        model2+=[nn.ReLU(True),]
+        model2+=[norm_layer(128),]
+
+        model3=[nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model3+=[nn.ReLU(True),]
+        model3+=[nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model3+=[nn.ReLU(True),]
+        model3+=[nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1, bias=True),]
+        model3+=[nn.ReLU(True),]
+        model3+=[norm_layer(256),]
+
+        model4=[nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model4+=[nn.ReLU(True),]
+        model4+=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model4+=[nn.ReLU(True),]
+        model4+=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model4+=[nn.ReLU(True),]
+        model4+=[norm_layer(512),]
+
+        model5=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model5+=[nn.ReLU(True),]
+        model5+=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model5+=[nn.ReLU(True),]
+        model5+=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model5+=[nn.ReLU(True),]
+        model5+=[norm_layer(512),]
+
+        model6=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model6+=[nn.ReLU(True),]
+        model6+=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model6+=[nn.ReLU(True),]
+        model6+=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model6+=[nn.ReLU(True),]
+        model6+=[norm_layer(512),]
+
+        model7=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model7+=[nn.ReLU(True),]
+        model7+=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model7+=[nn.ReLU(True),]
+        model7+=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model7+=[nn.ReLU(True),]
+        model7+=[norm_layer(512),]
+
+        model8=[nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=True),]
+        model8+=[nn.ReLU(True),]
+        model8+=[nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model8+=[nn.ReLU(True),]
+        model8+=[nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model8+=[nn.ReLU(True),]
+
+        model8+=[nn.Conv2d(256, 313, kernel_size=1, stride=1, padding=0, bias=True),]
+
+        self.model1 = nn.Sequential(*model1)
+        self.model2 = nn.Sequential(*model2)
+        self.model3 = nn.Sequential(*model3)
+        self.model4 = nn.Sequential(*model4)
+        self.model5 = nn.Sequential(*model5)
+        self.model6 = nn.Sequential(*model6)
+        self.model7 = nn.Sequential(*model7)
+        self.model8 = nn.Sequential(*model8)
+
+        self.softmax = nn.Softmax(dim=1)
+        self.model_out = nn.Conv2d(313, 2, kernel_size=1, padding=0, dilation=1, stride=1, bias=False)
+        self.upsample4 = nn.Upsample(scale_factor=4, mode='bilinear')
+
+    def forward(self, input_l):
+        conv1_2 = self.model1(self.normalize_l(input_l))
+        conv2_2 = self.model2(conv1_2)
+        conv3_3 = self.model3(conv2_2)
+        conv4_3 = self.model4(conv3_3)
+        conv5_3 = self.model5(conv4_3)
+        conv6_3 = self.model6(conv5_3)
+        conv7_3 = self.model7(conv6_3)
+        conv8_3 = self.model8(conv7_3)
+        out_reg = self.model_out(self.softmax(conv8_3))
+
+        return self.unnormalize_ab(self.upsample4(out_reg))
+
+def eccv16(pretrained=True):
+	model = ECCVGenerator()
+	if(pretrained):
+		import torch.utils.model_zoo as model_zoo
+		model.load_state_dict(model_zoo.load_url('https://colorizers.s3.us-east-2.amazonaws.com/colorization_release_v2-9b330a0b.pth',map_location='cpu',check_hash=True))
+	return model

app/colorizers/siggraph17.py

@@ -0,0 +1,168 @@
+import torch
+import torch.nn as nn
+
+from .base_color import *
+
+class SIGGRAPHGenerator(BaseColor):
+    def __init__(self, norm_layer=nn.BatchNorm2d, classes=529):
+        super(SIGGRAPHGenerator, self).__init__()
+
+        # Conv1
+        model1=[nn.Conv2d(4, 64, kernel_size=3, stride=1, padding=1, bias=True),]
+        model1+=[nn.ReLU(True),]
+        model1+=[nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True),]
+        model1+=[nn.ReLU(True),]
+        model1+=[norm_layer(64),]
+        # add a subsampling operation
+
+        # Conv2
+        model2=[nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=True),]
+        model2+=[nn.ReLU(True),]
+        model2+=[nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1, bias=True),]
+        model2+=[nn.ReLU(True),]
+        model2+=[norm_layer(128),]
+        # add a subsampling layer operation
+
+        # Conv3
+        model3=[nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model3+=[nn.ReLU(True),]
+        model3+=[nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model3+=[nn.ReLU(True),]
+        model3+=[nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model3+=[nn.ReLU(True),]
+        model3+=[norm_layer(256),]
+        # add a subsampling layer operation
+
+        # Conv4
+        model4=[nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model4+=[nn.ReLU(True),]
+        model4+=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model4+=[nn.ReLU(True),]
+        model4+=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model4+=[nn.ReLU(True),]
+        model4+=[norm_layer(512),]
+
+        # Conv5
+        model5=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model5+=[nn.ReLU(True),]
+        model5+=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model5+=[nn.ReLU(True),]
+        model5+=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model5+=[nn.ReLU(True),]
+        model5+=[norm_layer(512),]
+
+        # Conv6
+        model6=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model6+=[nn.ReLU(True),]
+        model6+=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model6+=[nn.ReLU(True),]
+        model6+=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model6+=[nn.ReLU(True),]
+        model6+=[norm_layer(512),]
+
+        # Conv7
+        model7=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model7+=[nn.ReLU(True),]
+        model7+=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model7+=[nn.ReLU(True),]
+        model7+=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model7+=[nn.ReLU(True),]
+        model7+=[norm_layer(512),]
+
+        # Conv7
+        model8up=[nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=True)]
+        model3short8=[nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+
+        model8=[nn.ReLU(True),]
+        model8+=[nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model8+=[nn.ReLU(True),]
+        model8+=[nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model8+=[nn.ReLU(True),]
+        model8+=[norm_layer(256),]
+
+        # Conv9
+        model9up=[nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=True),]
+        model2short9=[nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1, bias=True),]
+        # add the two feature maps above        
+
+        model9=[nn.ReLU(True),]
+        model9+=[nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1, bias=True),]
+        model9+=[nn.ReLU(True),]
+        model9+=[norm_layer(128),]
+
+        # Conv10
+        model10up=[nn.ConvTranspose2d(128, 128, kernel_size=4, stride=2, padding=1, bias=True),]
+        model1short10=[nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=True),]
+        # add the two feature maps above
+
+        model10=[nn.ReLU(True),]
+        model10+=[nn.Conv2d(128, 128, kernel_size=3, dilation=1, stride=1, padding=1, bias=True),]
+        model10+=[nn.LeakyReLU(negative_slope=.2),]
+
+        # classification output
+        model_class=[nn.Conv2d(256, classes, kernel_size=1, padding=0, dilation=1, stride=1, bias=True),]
+
+        # regression output
+        model_out=[nn.Conv2d(128, 2, kernel_size=1, padding=0, dilation=1, stride=1, bias=True),]
+        model_out+=[nn.Tanh()]
+
+        self.model1 = nn.Sequential(*model1)
+        self.model2 = nn.Sequential(*model2)
+        self.model3 = nn.Sequential(*model3)
+        self.model4 = nn.Sequential(*model4)
+        self.model5 = nn.Sequential(*model5)
+        self.model6 = nn.Sequential(*model6)
+        self.model7 = nn.Sequential(*model7)
+        self.model8up = nn.Sequential(*model8up)
+        self.model8 = nn.Sequential(*model8)
+        self.model9up = nn.Sequential(*model9up)
+        self.model9 = nn.Sequential(*model9)
+        self.model10up = nn.Sequential(*model10up)
+        self.model10 = nn.Sequential(*model10)
+        self.model3short8 = nn.Sequential(*model3short8)
+        self.model2short9 = nn.Sequential(*model2short9)
+        self.model1short10 = nn.Sequential(*model1short10)
+
+        self.model_class = nn.Sequential(*model_class)
+        self.model_out = nn.Sequential(*model_out)
+
+        self.upsample4 = nn.Sequential(*[nn.Upsample(scale_factor=4, mode='bilinear'),])
+        self.softmax = nn.Sequential(*[nn.Softmax(dim=1),])
+
+    def forward(self, input_A, input_B=None, mask_B=None):
+        if(input_B is None):
+            input_B = torch.cat((input_A*0, input_A*0), dim=1)
+        if(mask_B is None):
+            mask_B = input_A*0
+
+        conv1_2 = self.model1(torch.cat((self.normalize_l(input_A),self.normalize_ab(input_B),mask_B),dim=1))
+        conv2_2 = self.model2(conv1_2[:,:,::2,::2])
+        conv3_3 = self.model3(conv2_2[:,:,::2,::2])
+        conv4_3 = self.model4(conv3_3[:,:,::2,::2])
+        conv5_3 = self.model5(conv4_3)
+        conv6_3 = self.model6(conv5_3)
+        conv7_3 = self.model7(conv6_3)
+
+        conv8_up = self.model8up(conv7_3) + self.model3short8(conv3_3)
+        conv8_3 = self.model8(conv8_up)
+        conv9_up = self.model9up(conv8_3) + self.model2short9(conv2_2)
+        conv9_3 = self.model9(conv9_up)
+        conv10_up = self.model10up(conv9_3) + self.model1short10(conv1_2)
+        conv10_2 = self.model10(conv10_up)
+        out_reg = self.model_out(conv10_2)
+
+        conv9_up = self.model9up(conv8_3) + self.model2short9(conv2_2)
+        conv9_3 = self.model9(conv9_up)
+        conv10_up = self.model10up(conv9_3) + self.model1short10(conv1_2)
+        conv10_2 = self.model10(conv10_up)
+        out_reg = self.model_out(conv10_2)
+
+        return self.unnormalize_ab(out_reg)
+
+def siggraph17(pretrained=True):
+    model = SIGGRAPHGenerator()
+    if(pretrained):
+        import torch.utils.model_zoo as model_zoo
+        model.load_state_dict(model_zoo.load_url('https://colorizers.s3.us-east-2.amazonaws.com/siggraph17-df00044c.pth',map_location='cpu',check_hash=True))
+    return model
+

app/colorizers/util.py

@@ -0,0 +1,47 @@
+
+from PIL import Image
+import numpy as np
+from skimage import color
+import torch
+import torch.nn.functional as F
+from IPython import embed
+
+def load_img(img_path):
+	out_np = np.asarray(Image.open(img_path))
+	if(out_np.ndim==2):
+		out_np = np.tile(out_np[:,:,None],3)
+	return out_np
+
+def resize_img(img, HW=(256,256), resample=3):
+	return np.asarray(Image.fromarray(img).resize((HW[1],HW[0]), resample=resample))
+
+def preprocess_img(img_rgb_orig, HW=(256,256), resample=3):
+	# return original size L and resized L as torch Tensors
+	img_rgb_rs = resize_img(img_rgb_orig, HW=HW, resample=resample)
+	
+	img_lab_orig = color.rgb2lab(img_rgb_orig)
+	img_lab_rs = color.rgb2lab(img_rgb_rs)
+
+	img_l_orig = img_lab_orig[:,:,0]
+	img_l_rs = img_lab_rs[:,:,0]
+
+	tens_orig_l = torch.Tensor(img_l_orig)[None,None,:,:]
+	tens_rs_l = torch.Tensor(img_l_rs)[None,None,:,:]
+
+	return (tens_orig_l, tens_rs_l)
+
+def postprocess_tens(tens_orig_l, out_ab, mode='bilinear'):
+	# tens_orig_l 	1 x 1 x H_orig x W_orig
+	# out_ab 		1 x 2 x H x W
+
+	HW_orig = tens_orig_l.shape[2:]
+	HW = out_ab.shape[2:]
+
+	# call resize function if needed
+	if(HW_orig[0]!=HW[0] or HW_orig[1]!=HW[1]):
+		out_ab_orig = F.interpolate(out_ab, size=HW_orig, mode='bilinear')
+	else:
+		out_ab_orig = out_ab
+
+	out_lab_orig = torch.cat((tens_orig_l, out_ab_orig), dim=1)
+	return color.lab2rgb(out_lab_orig.data.cpu().numpy()[0,...].transpose((1,2,0)))

app/config.py

@@ -46,7 +46,8 @@ class Settings(BaseSettings):
         "Colorized using GAN-Colorization-Model"
     )
     INFERENCE_PROVIDER: str = os.getenv("INFERENCE_PROVIDER", "fal-ai")
-    INFERENCE_MODEL: str = os.getenv("INFERENCE_MODEL", "black-forest-labs/FLUX.1-Kontext-dev")
+    # Note: black-forest-labs interface not used in main.py - only used in main_sdxl.py
+    INFERENCE_MODEL: str = os.getenv("INFERENCE_MODEL", "")
     INFERENCE_TIMEOUT: int = int(os.getenv("INFERENCE_TIMEOUT", "180"))
     HF_TOKEN: str = os.getenv("HF_TOKEN", "")
     

app/main.py

@@ -6,9 +6,11 @@ import uuid
 import os
 import io
 import json
+import logging
 from PIL import Image
 import torch
 from torchvision import transforms
+import numpy as np
 from app.database import (
     get_database,
     log_api_call,
@@ -22,6 +24,17 @@ try:
 except ImportError:
     firebase_auth = None
 
+# Import CCO colorizers
+try:
+    from app.colorizers import eccv16, siggraph17
+    from app.colorizers.util import preprocess_img, postprocess_tens
+    CCO_AVAILABLE = True
+except ImportError as e:
+    print(f"⚠️ CCO colorizers not available: {e}")
+    CCO_AVAILABLE = False
+
+logger = logging.getLogger(__name__)
+
 # -------------------------------------------------
 # 🚀 FastAPI App
 # -------------------------------------------------
@@ -63,10 +76,10 @@ MEDIA_CLICK_DEFAULT_CATEGORY = os.getenv("DEFAULT_CATEGORY_FALLBACK", "69368fcd2
 MODEL_REPO = "Hammad712/GAN-Colorization-Model"
 MODEL_FILENAME = "generator.pt"
 
-print("⬇️ Downloading model...")
+print("⬇️ Downloading GAN model...")
 model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILENAME)
 
-print("📦 Loading model weights...")
+print("📦 Loading GAN model weights...")
 state_dict = torch.load(model_path, map_location="cpu")
 
 # NOTE: Replace with real model architecture
@@ -75,14 +88,76 @@ state_dict = torch.load(model_path, map_location="cpu")
 # model.load_state_dict(state_dict)
 # model.eval()
 
-def colorize_image(img: Image.Image):
-    """ Dummy colorizer (replace with real model.predict) """
+# -------------------------------------------------
+# 🧠 Load CCO Colorization Models
+# -------------------------------------------------
+cco_models = {}
+if CCO_AVAILABLE:
+    print("📦 Loading CCO models...")
+    try:
+        cco_models["eccv16"] = eccv16(pretrained=True).eval()
+        cco_models["siggraph17"] = siggraph17(pretrained=True).eval()
+        print("✅ CCO models loaded successfully!")
+    except Exception as e:
+        print(f"⚠️ Failed to load CCO models: {e}")
+        CCO_AVAILABLE = False
+
+def colorize_image_gan(img: Image.Image):
+    """ GAN colorizer (dummy implementation - replace with real model.predict) """
     transform = transforms.ToTensor()
     tensor = transform(img.convert("L")).unsqueeze(0)
     tensor = tensor.repeat(1, 3, 1, 1)
     output_img = transforms.ToPILImage()(tensor.squeeze())
     return output_img
 
+def colorize_image_cco(img: Image.Image, model_name: str = "eccv16"):
+    """ CCO colorizer using eccv16 or siggraph17 model """
+    if not CCO_AVAILABLE:
+        raise ValueError("CCO models are not available")
+    
+    if model_name not in ["eccv16", "siggraph17"]:
+        model_name = "eccv16"  # Default to eccv16
+    
+    model = cco_models.get(model_name)
+    if model is None:
+        raise ValueError(f"CCO model '{model_name}' not loaded")
+    
+    # Convert PIL Image to numpy array
+    oimg = np.asarray(img)
+    if oimg.ndim == 2:
+        oimg = np.tile(oimg[:,:,None], 3)
+    
+    # Preprocess image
+    (tens_l_orig, tens_l_rs) = preprocess_img(oimg)
+    
+    # Run model inference
+    with torch.no_grad():
+        out_ab = model(tens_l_rs)
+    
+    # Postprocess output
+    output_rgb = postprocess_tens(tens_l_orig, out_ab)
+    
+    # Convert numpy array back to PIL Image
+    output_img = Image.fromarray((output_rgb * 255).astype(np.uint8))
+    return output_img
+
+def colorize_image(img: Image.Image, model_type: str = "gan", cco_model: str = "eccv16"):
+    """
+    Colorize image using specified model
+    
+    Args:
+        img: PIL Image to colorize
+        model_type: "gan" or "cco"
+        cco_model: "eccv16" or "siggraph17" (only used if model_type is "cco")
+    
+    Returns:
+        Colorized PIL Image
+    """
+    if model_type == "cco":
+        return colorize_image_cco(img, cco_model)
+    else:
+        return colorize_image_gan(img)
+
 # -------------------------------------------------
 # 🗄️ MongoDB Initialization
 # -------------------------------------------------
@@ -223,6 +298,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"
 ):
     import time
     start_time = time.time()
@@ -237,6 +313,50 @@ async def colorize(
     if not effective_category_id:
         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
+    
+    if model:
+        model = model.strip().lower()
+        if model == "cco" or model.startswith("cco-"):
+            if not CCO_AVAILABLE:
+                error_msg = "CCO models are not available"
+                log_api_call(
+                    endpoint="/colorize",
+                    method="POST",
+                    status_code=400,
+                    error=error_msg,
+                    ip_address=ip_address
+                )
+                log_colorization(
+                    result_id=None,
+                    model_type="cco",
+                    processing_time=None,
+                    user_id=effective_user_id,
+                    ip_address=ip_address,
+                    status="failed",
+                    error=error_msg
+                )
+                raise HTTPException(status_code=400, detail=error_msg)
+            
+            model_type = "cco"
+            if model == "cco-eccv16":
+                cco_model = "eccv16"
+                model_type_for_log = "cco-eccv16"
+            elif model == "cco-siggraph17":
+                cco_model = "siggraph17"
+                model_type_for_log = "cco-siggraph17"
+            else:
+                # Default to eccv16 if just "cco" is specified
+                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"
+
     if not file.content_type.startswith("image/"):
         error_msg = "Invalid file type"
         log_api_call(
@@ -249,7 +369,7 @@ async def colorize(
         # Log failed colorization
         log_colorization(
             result_id=None,
-            model_type="gan",
+            model_type=model_type_for_log,
             processing_time=None,
             user_id=effective_user_id,
             ip_address=ip_address,
@@ -260,7 +380,7 @@ async def colorize(
 
     try:
         img = Image.open(io.BytesIO(await file.read()))
-        output_img = colorize_image(img)
+        output_img = colorize_image(img, model_type=model_type, cco_model=cco_model)
 
         processing_time = time.time() - start_time
 
@@ -276,13 +396,14 @@ async def colorize(
             "success": True,
             "result_id": result_id_clean,
             "download_url": f"{base_url}/results/{result_id}",
-            "api_download": f"{base_url}/download/{result_id_clean}"
+            "api_download": f"{base_url}/download/{result_id_clean}",
+            "model_used": model_type_for_log
         }
         
         # Log to MongoDB (colorization_db -> colorizations)
         log_colorization(
             result_id=result_id_clean,
-            model_type="gan",
+            model_type=model_type_for_log,
             processing_time=processing_time,
             user_id=effective_user_id,
             ip_address=ip_address,
@@ -293,7 +414,7 @@ async def colorize(
             endpoint="/colorize",
             method="POST",
             status_code=200,
-            request_data={"filename": file.filename, "content_type": file.content_type},
+            request_data={"filename": file.filename, "content_type": file.content_type, "model": model},
             response_data=response_data,
             user_id=effective_user_id,
             ip_address=ip_address
@@ -314,7 +435,7 @@ async def colorize(
         # Log failed colorization to colorizations collection
         log_colorization(
             result_id=None,
-            model_type="gan",
+            model_type=model_type_for_log,
             processing_time=None,
             user_id=effective_user_id,
             ip_address=ip_address,

requirements.txt

@@ -17,4 +17,5 @@ safetensors
 ftfy
 httpx
 email-validator
-pymongo
+pymongo
+scikit-image