Spaces:
Running on Zero
Running on Zero
File size: 7,652 Bytes
f0d9a3e 4bb998d f0d9a3e 4bb998d f0d9a3e 4bb998d f0d9a3e 4bb998d f0d9a3e 4bb998d f0d9a3e 4bb998d f0d9a3e 4bb998d f0d9a3e 4bb998d f0d9a3e 4bb998d 49d3665 f0d9a3e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 | """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
|