lib/image_utils.py

from PIL import Image
import io
from pathlib import Path
import time
import traceback
from typing import Union, Optional
from .api_utils import upload_to_freeimage

def prepare_image(image: Union[str, bytes, Image.Image, None], status_tracker) -> Optional[str]:
    """
    Prepare an image for use with LumaAI by resizing and uploading to CDN.
    
    Args:
        image: Input image (can be path, bytes, or PIL Image)
        status_tracker: StatusTracker instance for progress updates
        
    Returns:
        Optional[str]: CDN URL of the prepared image, or None if no image provided
        
    Raises:
        Exception: If image preparation fails
    """
    if image is None:
        return None
        
    try:
        status_tracker.update_message("Preparing your image for its big moment...", 0.01)
        
        # Convert to PIL Image if needed
        if isinstance(image, str):
            image = Image.open(image)
        elif isinstance(image, bytes):
            image = Image.open(io.BytesIO(image))
        elif not isinstance(image, Image.Image):
            raise Exception("That doesn't look like an image (unless I need glasses)")

        # Resize image to 512x512
        image = image.resize((512, 512), Image.Resampling.LANCZOS)
        status_tracker.add_step("Image resized to 512x512", 0.02)

        # Convert to RGB if necessary
        if image.mode not in ('RGB', 'RGBA'):
            image = image.convert('RGB')

        # Save to temporary file
        temp_dir = Path("temp")
        temp_dir.mkdir(exist_ok=True)
        temp_path = temp_dir / f"temp_image_{int(time.time())}.png"
        image.save(str(temp_path), format='PNG', optimize=True)
        
        # Upload to freeimage.host
        try:
            cdn_url = upload_to_freeimage(temp_path, status_tracker)
            
            # Clean up temporary file
            if temp_path.exists():
                temp_path.unlink()
                
            return cdn_url
        except Exception as e:
            # Clean up temporary file in case of error
            if temp_path.exists():
                temp_path.unlink()
            raise Exception(f"Image upload failed: {str(e)}")
            
    except Exception as e:
        print(f"Error preparing image: {str(e)}")
        print(traceback.format_exc())
        raise Exception(f"Image preparation failed: {str(e)}")