app/colorize_model.py

"""
ColorizeNet model wrapper for image colorization
"""
import logging
import torch
import numpy as np
from PIL import Image
import cv2
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, StableDiffusionXLControlNetPipeline
from diffusers.utils import load_image
from transformers import pipeline
from huggingface_hub import hf_hub_download

logger = logging.getLogger(__name__)

class ColorizeModel:
    """Wrapper for ColorizeNet model"""
    
    def __init__(self, model_id: str = "rsortino/ColorizeNet"):
        """
        Initialize the ColorizeNet model
        
        Args:
            model_id: Hugging Face model ID for ColorizeNet
        """
        self.model_id = model_id
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        logger.info("Using device: %s", self.device)
        self.dtype = torch.float16 if self.device == "cuda" else torch.float32
        
        try:
            # Try loading as ControlNet with Stable Diffusion
            logger.info("Attempting to load ColorizeNet as ControlNet...")
            try:
                # Load ControlNet model
                self.controlnet = ControlNetModel.from_pretrained(
                    model_id,
                    torch_dtype=self.dtype
                )
                
                # Try SDXL first, fallback to SD 1.5
                try:
                    self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
                        "stabilityai/stable-diffusion-xl-base-1.0",
                        controlnet=self.controlnet,
                        torch_dtype=self.dtype,
                        safety_checker=None,
                        requires_safety_checker=False
                    )
                    logger.info("Loaded with SDXL base model")
                except:
                    self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
                        "runwayml/stable-diffusion-v1-5",
                        controlnet=self.controlnet,
                        torch_dtype=self.dtype,
                        safety_checker=None,
                        requires_safety_checker=False
                    )
                    logger.info("Loaded with SD 1.5 base model")
                
                self.pipe.to(self.device)
                
                # Enable memory efficient attention if available
                if hasattr(self.pipe, "enable_xformers_memory_efficient_attention"):
                    try:
                        self.pipe.enable_xformers_memory_efficient_attention()
                        logger.info("XFormers memory efficient attention enabled")
                    except Exception as e:
                        logger.warning("Could not enable XFormers: %s", str(e))
                
                logger.info("ColorizeNet model loaded successfully as ControlNet")
                self.model_type = "controlnet"
                
            except Exception as e:
                logger.warning("Failed to load as ControlNet: %s", str(e))
                # Fallback: try as image-to-image pipeline
                logger.info("Trying to load as image-to-image pipeline...")
                self.pipe = pipeline(
                    "image-to-image",
                    model=model_id,
                    device=0 if self.device == "cuda" else -1,
                    torch_dtype=self.dtype
                )
                logger.info("ColorizeNet model loaded using image-to-image pipeline")
                self.model_type = "pipeline"
                
        except Exception as e:
            logger.error("Failed to load ColorizeNet model: %s", str(e))
            raise RuntimeError(f"Could not load ColorizeNet model: {str(e)}")
    
    def preprocess_image(self, image: Image.Image) -> Image.Image:
        """
        Preprocess image for colorization
        
        Args:
            image: PIL Image
            
        Returns:
            Preprocessed PIL Image
        """
        # Convert to grayscale if needed
        if image.mode != "L":
            # Convert to grayscale
            image = image.convert("L")
        
        # Convert back to RGB (grayscale image with 3 channels)
        image = image.convert("RGB")
        
        # Resize to standard size (512x512 for SD models)
        image = image.resize((512, 512), Image.Resampling.LANCZOS)
        
        return image
    
    def colorize(self, image: Image.Image, num_inference_steps: int = 20) -> Image.Image:
        """
        Colorize a grayscale image
        
        Args:
            image: PIL Image (grayscale or color)
            num_inference_steps: Number of inference steps
            
        Returns:
            Colorized PIL Image
        """
        try:
            # Preprocess image
            control_image = self.preprocess_image(image)
            original_size = image.size
            
            # Prepare prompt for colorization
            prompt = "colorize this black and white image, high quality, detailed, vibrant colors, natural colors"
            negative_prompt = "black and white, grayscale, monochrome, low quality, blurry, desaturated"
            
            # Generate colorized image based on model type
            if self.model_type == "controlnet":
                # Use ControlNet pipeline
                result = self.pipe(
                    prompt=prompt,
                    image=control_image,
                    negative_prompt=negative_prompt,
                    num_inference_steps=num_inference_steps,
                    guidance_scale=7.5,
                    controlnet_conditioning_scale=1.0,
                    generator=torch.Generator(device=self.device).manual_seed(42)
                )
                
                if isinstance(result, dict) and "images" in result:
                    colorized = result["images"][0]
                elif isinstance(result, list) and len(result) > 0:
                    colorized = result[0]
                else:
                    colorized = result
            else:
                # Use pipeline directly
                result = self.pipe(
                    control_image,
                    prompt=prompt,
                    num_inference_steps=num_inference_steps
                )
                
                if isinstance(result, dict) and "images" in result:
                    colorized = result["images"][0]
                elif isinstance(result, list) and len(result) > 0:
                    colorized = result[0]
                else:
                    colorized = result
            
            # Ensure we have a PIL Image
            if not isinstance(colorized, Image.Image):
                if isinstance(colorized, np.ndarray):
                    # Handle numpy array
                    if colorized.dtype != np.uint8:
                        colorized = (colorized * 255).astype(np.uint8)
                    if len(colorized.shape) == 3 and colorized.shape[2] == 3:
                        colorized = Image.fromarray(colorized, 'RGB')
                    else:
                        colorized = Image.fromarray(colorized)
                elif torch.is_tensor(colorized):
                    # Handle torch tensor
                    colorized = colorized.cpu().permute(1, 2, 0).numpy()
                    colorized = (colorized * 255).astype(np.uint8)
                    colorized = Image.fromarray(colorized, 'RGB')
                else:
                    raise ValueError(f"Unexpected output type: {type(colorized)}")
            
            # Resize back to original size
            if original_size != (512, 512):
                colorized = colorized.resize(original_size, Image.Resampling.LANCZOS)
            
            return colorized
            
        except Exception as e:
            logger.error("Error during colorization: %s", str(e))
            raise