Update generation.py

generation.py

@@ -5,76 +5,101 @@ import re
 from datetime import datetime
 from typing import Optional
 
+# Configuración de carpetas
 OUTPUT_DIR = "generated_images"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 def generate_image_from_prompt(
     prompt: str,
     negative_prompt: str = "",
-    model_name: str = "ignored", # Ignoramos lo que venga de fuera, usaremos el bueno aquí
+    model_name: str = "ignored", # Este argumento lo ignoramos para usar los hardcodeados seguros
     seed: Optional[int] = None,
 ) -> tuple[Optional[str], str]:
-    try:
-        # 1. Limpieza de Clave (Seguridad contra espacios invisibles)
-        api_key = os.getenv("OPENROUTER_API_KEY")
-        if not api_key:
-            return None, "❌ Falta OPENROUTER_API_KEY en los Secrets."
-        api_key = api_key.strip()
 
-        # 2. Configuración del modelo real de OpenRouter
-        # Flux Schnell es muy barato y rápido
-        real_model = "black-forest-labs/flux-schnell"
+    # 1. VALIDACIÓN DE CREDENCIALES
+    api_key = os.getenv("OPENROUTER_API_KEY")
+    if not api_key:
+        return None, "❌ Error Crítico: No existe OPENROUTER_API_KEY en Secrets."
+    
+    api_key = api_key.strip() # Limpieza de seguridad
 
-        print(f"Enviando prompt a OpenRouter: {prompt}")
+    # 2. DEFINICIÓN DE MODELOS (Principal y Respaldo)
+    # El ID correcto verificado es con '-1-'
+    primary_model = "black-forest-labs/flux-1-schnell"
+    backup_model = "stabilityai/stable-diffusion-xl-base-1.0"
 
-        response = requests.post(
-            url="https://openrouter.ai/api/v1/chat/completions",
-            headers={
+    models_to_try = [primary_model, backup_model]
+
+    last_error = ""
+
+    # 3. BUCLE DE INTENTOS
+    for model in models_to_try:
+        try:
+            print(f"🔄 Intentando generar con modelo: {model}...")
+            
+            headers = {
                 "Authorization": f"Bearer {api_key}",
                 "Content-Type": "application/json",
-                "HTTP-Referer": "https://huggingface.co", # Requerido por OpenRouter
+                "HTTP-Referer": "https://huggingface.co",
                 "X-Title": "Sofia AI Studio",
-            },
-            json={
-                "model": real_model,
+            }
+
+            payload = {
+                "model": model,
                 "messages": [
                     {"role": "user", "content": prompt}
                 ]
             }
-        )
-        
-        # 3. Procesar respuesta
-        if response.status_code != 200:
-            return None, f"❌ Error OpenRouter ({response.status_code}): {response.text}"
 
-        data = response.json()
-        
-        # OpenRouter con Flux suele devolver la URL dentro del contenido del mensaje
-        # A veces viene como markdown: ![image](url)
-        content = data['choices'][0]['message']['content']
-        
-        # Buscar la URL de la imagen usando una expresión regular
-        url_match = re.search(r'\((https://.*?)\)', content) # Busca (https://...)
-        if not url_match:
-             # Si no está entre paréntesis, buscamos directamente http
-             url_match = re.search(r'(https://[^\s]+)', content)
-             
-        if url_match:
-            image_url = url_match.group(1)
-            # Descargar la imagen al disco
-            img_data = requests.get(image_url).content
-            
-            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-            if seed is None: seed = random.randint(0, 9999)
-            filename = f"sofia_{timestamp}_{seed}.png"
-            file_path = os.path.join(OUTPUT_DIR, filename)
+            response = requests.post(
+                url="https://openrouter.ai/api/v1/chat/completions",
+                headers=headers,
+                json=payload,
+                timeout=45 # Timeout para evitar bloqueos eternos
+            )
+
+            # Si hay error 400/500, pasamos al siguiente modelo
+            if response.status_code != 200:
+                error_detail = response.text
+                print(f"⚠️ Fallo con {model}: {error_detail}")
+                last_error = f"Error {response.status_code} en {model}: {error_detail}"
+                continue # Salta al siguiente modelo del bucle
+
+            # Si es 200 OK, procesamos la imagen
+            data = response.json()
             
-            with open(file_path, 'wb') as f:
-                f.write(img_data)
+            # OpenRouter devuelve la imagen dentro del contenido del mensaje (Markdown o URL directa)
+            if "choices" in data and len(data["choices"]) > 0:
+                content = data['choices'][0]['message']['content']
                 
-            return file_path, f"✅ ¡ÉXITO! Imagen generada con {real_model}"
-        else:
-            return None, f"⚠️ OpenRouter respondió texto pero no encontré imagen: {content}"
+                # Buscamos la URL con regex (formatos markdown ![img](url) o url directa)
+                url_match = re.search(r'\((https://.*?)\)', content)
+                if not url_match:
+                    url_match = re.search(r'(https://[^\s]+\.(png|jpg|jpeg|webp))', content)
+                
+                if url_match:
+                    image_url = url_match.group(1)
+                    
+                    # Descargamos la imagen
+                    img_data = requests.get(image_url).content
+                    
+                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                    if seed is None: seed = random.randint(0, 9999)
+                    filename = f"sofia_{timestamp}_{seed}.png"
+                    file_path = os.path.join(OUTPUT_DIR, filename)
+                    
+                    with open(file_path, 'wb') as f:
+                        f.write(img_data)
+                    
+                    return file_path, f"✅ ÉXITO: Imagen creada con {model}"
+                else:
+                    last_error = f"La API respondió texto pero no vi imagen: {content[:50]}..."
+            else:
+                last_error = f"Respuesta vacía o formato desconocido: {data}"
+
+        except Exception as e:
+            last_error = f"Excepción técnica con {model}: {str(e)}"
+            continue
 
-    except Exception as e:
-        return None, f"❌ Error técnico crítico: {str(e)}"
+    # Si llega aquí, fallaron todos los modelos
+    return None, f"❌ ERROR FATAL: Fallaron todos los intentos.\nÚltimo error: {last_error}"