optimize FLUX model for memory and speed - reduce resolution and steps

app.py

@@ -191,7 +191,7 @@ def load_image_model(model_name):
                     print("🚀 Cargando FLUX Pipeline...")
                     pipe = FluxPipeline.from_pretrained(
                         model_name,
-                        torch_dtype=torch.bfloat16,
+                        torch_dtype=torch.float32,  # ✅ Cambiado a float32 para ahorrar memoria
                         use_auth_token=HF_TOKEN if HF_TOKEN else None
                     )
                     # Solo usar enable_model_cpu_offload si hay acelerador disponible
@@ -201,6 +201,12 @@ def load_image_model(model_name):
                     except Exception as offload_error:
                         print(f"⚠️ No se pudo habilitar CPU offload: {offload_error}")
                         print("✅ FLUX cargado sin CPU offload")
+                    
+                    # Optimizaciones adicionales para ahorrar memoria
+                    pipe.enable_attention_slicing()
+                    pipe.enable_vae_slicing()
+                    print("✅ Optimizaciones de memoria aplicadas a FLUX")
+                    
                 except Exception as e:
                     print(f"❌ Error cargando FLUX: {e}")
                     # Fallback a Stable Diffusion
@@ -420,7 +426,7 @@ def generate_text(prompt, model_name, max_length=100):
         return f"Error generando texto: {str(e)}"
 
 def generate_image(prompt, model_name, num_inference_steps=20):
-    """Generar imagen con el modelo seleccionado - versión simplificada con soporte para FLUX"""
+    """Generar imagen con el modelo seleccionado - versión optimizada para Spaces"""
     try:
         print(f"Generando imagen con modelo: {model_name}")
         print(f"Prompt: {prompt}")
@@ -437,21 +443,22 @@ def generate_image(prompt, model_name, num_inference_steps=20):
         model_data = load_image_model(model_name)
         pipeline = model_data["pipeline"]
         
-        # Configuración específica para FLUX
+        # Configuración específica para FLUX - OPTIMIZADA para Spaces
         if "flux" in model_name.lower():
             import random
             # Generar un seed aleatorio para cada imagen
             random_seed = random.randint(0, 999999)
             print(f"🎲 Usando seed aleatorio para FLUX: {random_seed}")
-            print(f"🔧 Parámetros FLUX: guidance_scale=3.5, steps=50, max_seq=512")
+            print(f"🔧 Parámetros FLUX OPTIMIZADOS: guidance_scale=3.5, steps=25, max_seq=256, height=512")
             
+            # Parámetros optimizados para reducir memoria y tiempo
             image = pipeline(
                 prompt,
-                height=1024,
-                width=1024,
+                height=512,  # ✅ Reducido de 1024 a 512 para ahorrar memoria
+                width=512,   # ✅ Reducido de 1024 a 512 para ahorrar memoria
                 guidance_scale=3.5,  # ✅ Valor recomendado por la documentación
-                num_inference_steps=50,  # ✅ Valor recomendado por la documentación
-                max_sequence_length=512,  # ✅ Valor recomendado por la documentación
+                num_inference_steps=25,  # ✅ Reducido de 50 a 25 para ahorrar tiempo
+                max_sequence_length=256,  # ✅ Reducido de 512 a 256 para ahorrar memoria
                 generator=torch.Generator("cpu").manual_seed(random_seed)  # ✅ Seed aleatorio
             ).images[0]
         else:
@@ -462,11 +469,11 @@ def generate_image(prompt, model_name, num_inference_steps=20):
                 guidance_scale=7.5
         ).images[0]
         
-        print("Imagen generada exitosamente")
+        print("✅ Imagen generada exitosamente")
         return image
         
     except Exception as e:
-        print(f"Error generando imagen: {str(e)}")
+        print(f"❌ Error generando imagen: {str(e)}")
         return f"Error generando imagen: {str(e)}"
 
 def generate_video(prompt, model_name, num_frames=16, num_inference_steps=20):