src/model_handler.py

"""Model handler for Flux Seamless Texture LoRA."""

import logging
import os
import io
from typing import Optional, Dict, Any, List, Generator
from pathlib import Path
from PIL import Image
from huggingface_hub import InferenceClient

from config.settings import MODEL_ID, MODEL_PROVIDER, DEFAULT_PARAMS
from src.utils import generate_seed, validate_params
from src.image_processor import save_image

logger = logging.getLogger(__name__)


class ModelHandler:
    """Handler for managing the Flux LoRA model."""
    
    def __init__(self):
        """Initialize the model handler."""
        self.client = None
        self.model_id = self._normalize_model_id(MODEL_ID)
        self._init_client()
    
    @staticmethod
    def _normalize_model_id(model_id: str) -> str:
        """Normalize model id across Gradio/HF conventions."""
        # Some gr.load usages prefix with "models/". HF Inference expects repo id.
        if model_id.startswith("models/"):
            return model_id[len("models/") :]
        return model_id
    
    def _init_client(self) -> None:
        """Initialize HF Inference client (Spaces-friendly)."""
        try:
            token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
            # In HF Spaces, prefer secrets; token is optional for public models.
            self.client = InferenceClient(token=token)
            logger.info(f"InferenceClient initialized for model: {self.model_id}")
        except Exception as e:
            logger.error(f"Error initializing InferenceClient: {e}")
            raise
    
    def generate(
        self,
        prompt: str,
        negative_prompt: str = "",
        guidance_scale: float = 7.5,
        num_inference_steps: int = 50,
        seed: Optional[int] = None,
        width: int = 1024,
        height: int = 1024,
        cfg_scale: Optional[float] = None,
        lora_strength: float = 1.0,
    ) -> tuple[Any, Dict[str, Any]]:
        """Generate an image from a prompt.
        
        Args:
            prompt: Text prompt for generation.
            negative_prompt: Negative prompt.
            guidance_scale: Guidance scale for generation.
            num_inference_steps: Number of inference steps.
            seed: Random seed (None for random).
            width: Image width.
            height: Image height.
            cfg_scale: CFG scale (uses guidance_scale if None).
            lora_strength: LoRA strength.
            
        Returns:
            Tuple of (image, metadata_dict).
        """
        if not self.client:
            raise RuntimeError("Inference client not initialized")
        
        # Use cfg_scale if provided, otherwise use guidance_scale
        if cfg_scale is None:
            cfg_scale = guidance_scale
        
        # Generate seed if not provided
        if seed is None:
            seed = generate_seed()
        
        # Prepare parameters
        params = {
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "guidance_scale": guidance_scale,
            "num_inference_steps": num_inference_steps,
            "seed": seed,
            "width": width,
            "height": height,
            "cfg_scale": cfg_scale,
            "lora_strength": lora_strength,
        }
        
        # Validate parameters
        is_valid, error = validate_params(params)
        if not is_valid:
            raise ValueError(f"Invalid parameters: {error}")
        
        try:
            logger.info(f"Generating image with prompt: {prompt[:50]}...")
            
            # HF Inference parameters (best-effort; backend may ignore unknown keys).
            hf_params: Dict[str, Any] = {
                "negative_prompt": negative_prompt,
                "guidance_scale": float(guidance_scale),
                "num_inference_steps": int(num_inference_steps),
                "width": int(width),
                "height": int(height),
                "seed": int(seed),
                # Some backends accept "cfg_scale" or "lora_scale"; keep as best-effort.
                "cfg_scale": float(cfg_scale),
                "lora_scale": float(lora_strength),
            }
            
            image: Optional[Image.Image] = None
            
            # Preferred path: InferenceClient.text_to_image
            try:
                image = self.client.text_to_image(
                    prompt=prompt,
                    model=self.model_id,
                    **hf_params,
                )
            except TypeError:
                # Older signature: pass only supported keys
                image = self.client.text_to_image(
                    prompt=prompt,
                    model=self.model_id,
                )
            except Exception as e:
                # Fallback path: raw POST
                logger.warning(f"text_to_image failed ({e}); falling back to raw POST.")
                payload = {"inputs": prompt, "parameters": hf_params}
                raw = self.client.post(model=self.model_id, json=payload)
                image = Image.open(io.BytesIO(raw)).convert("RGB")
            
            if image is None:
                raise RuntimeError("No image returned from HF Inference")

            # IMPORTANT (Gradio/ORJSON): algumas imagens vindas de providers podem
            # carregar metadados com chaves não-string (ex: EXIF com ints). O Gradio
            # pode tentar serializar isso e explodir com "Dict key must be str".
            # Re-encode como PNG em memória e recarrega para "limpar" metadados.
            try:
                cleaned = io.BytesIO()
                image = image.convert("RGB")
                image.save(cleaned, format="PNG")
                cleaned.seek(0)
                image = Image.open(cleaned)
                image.load()
            except Exception as e:
                logger.warning(f"Could not sanitize image metadata: {e}")
            
            # Save image with metadata
            image_path = save_image(image, prompt, params)
            
            metadata = {
                "seed": seed,
                "image_path": str(image_path),
            }
            
            logger.info("Image generated successfully")
            return image, metadata
            
        except Exception as e:
            logger.error(f"Error generating image: {e}")
            raise
    
    def generate_batch(
        self,
        prompts: List[str],
        base_params: Optional[Dict[str, Any]] = None,
    ) -> Generator[tuple[Any, Dict[str, Any], int], None, None]:
        """Generate multiple images in batch.
        
        Args:
            prompts: List of prompts to generate.
            base_params: Base parameters to use for all generations.
            
        Yields:
            Tuple of (image, metadata, index) for each generation.
        """
        if base_params is None:
            base_params = {}
        
        total = len(prompts)
        logger.info(f"Starting batch generation of {total} images")
        
        for idx, prompt in enumerate(prompts):
            try:
                # Merge base params with defaults
                params = {**DEFAULT_PARAMS, **base_params}
                
                image, metadata = self.generate(
                    prompt=prompt,
                    **params
                )
                
                yield image, metadata, idx
                
            except Exception as e:
                logger.error(f"Error generating image {idx + 1}/{total}: {e}")
                yield None, {"error": str(e), "index": idx}, idx