Arreglar imágenes negras - Mejorar carga de modelos, manejo de variants fp16, verificación de imágenes generadas

app.py

@@ -73,9 +73,8 @@ MODELS = {
     "Midjourney Style (prompthero/openjourney)": "prompthero/openjourney",
     "Orange Mixs (WarriorMama777/OrangeMixs)": "WarriorMama777/OrangeMixs",
     "Kohaku V2.1 (KBlueLeaf/kohaku-v2.1)": "KBlueLeaf/kohaku-v2.1",
-    # Modelos avanzados que aprovechan H200
+    # Modelos avanzados que aprovechan H200 (solo los que existen)
     "SDXL Lightning (ByteDance/SDXL-Lightning)": "ByteDance/SDXL-Lightning",
-    "SDXL Lightning 4Step (ByteDance/SDXL-Lightning-4Step)": "ByteDance/SDXL-Lightning-4Step",
     "FLUX.1-Kontext-Dev (API External)": "api_external",
 }
 
@@ -84,8 +83,6 @@ if HF_TOKEN:
     FLUX_MODELS = {
         "FLUX.1-dev (black-forest-labs/FLUX.1-dev)": "black-forest-labs/FLUX.1-dev",
         "FLUX.1-schnell (black-forest-labs/FLUX.1-schnell)": "black-forest-labs/FLUX.1-schnell",
-        # Modelos FLUX adicionales que aprovechan H200
-        "FLUX.1-pro (black-forest-labs/FLUX.1-pro)": "black-forest-labs/FLUX.1-pro",
     }
     MODELS.update(FLUX_MODELS)
     print("🔓 Modelos FLUX habilitados con autenticación")
@@ -116,22 +113,38 @@ def load_model(model_id):
         try:
             start_time = time.time()
             
+            # Determinar si usar variant fp16 basado en el modelo
+            use_fp16_variant = False
+            if torch.cuda.is_available():
+                # Solo usar fp16 variant para modelos que lo soportan
+                fp16_supported_models = [
+                    "stabilityai/sdxl-turbo",
+                    "stabilityai/sd-turbo", 
+                    "stabilityai/stable-diffusion-xl-base-1.0",
+                    "runwayml/stable-diffusion-v1-5",
+                    "CompVis/stable-diffusion-v1-4"
+                ]
+                use_fp16_variant = any(model in model_id for model in fp16_supported_models)
+                print(f"🔧 FP16 variant: {'✅ Habilitado' if use_fp16_variant else '❌ Deshabilitado'} para {model_id}")
+            
             # Usar token de autenticación si está disponible
             if HF_TOKEN and ("flux" in model_id.lower() or "black-forest" in model_id.lower()):
                 print(f"🔐 Cargando modelo gated: {model_id}")
                 print(f"🔑 Usando token de autenticación...")
+                
+                # Para modelos FLUX, no usar variant fp16
                 pipe = DiffusionPipeline.from_pretrained(
                     model_id, 
                     torch_dtype=torch_dtype,
                     use_auth_token=HF_TOKEN,
-                    variant="fp16" if torch.cuda.is_available() else None
+                    variant="fp16" if use_fp16_variant else None
                 )
             else:
                 print(f"📦 Cargando modelo público: {model_id}")
                 pipe = DiffusionPipeline.from_pretrained(
                     model_id, 
                     torch_dtype=torch_dtype,
-                    variant="fp16" if torch.cuda.is_available() else None
+                    variant="fp16" if use_fp16_variant else None
                 )
             
             load_time = time.time() - start_time
@@ -144,34 +157,65 @@ def load_model(model_id):
             if torch.cuda.is_available():
                 print("🔧 Aplicando optimizaciones para H200...")
                 
-                # Habilitar optimizaciones de memoria
+                # Habilitar optimizaciones de memoria (más conservadoras)
                 if hasattr(pipe, 'enable_attention_slicing'):
                     pipe.enable_attention_slicing()
                     print("✅ Attention slicing habilitado")
                 
-                if hasattr(pipe, 'enable_model_cpu_offload'):
+                # Solo usar CPU offload para modelos grandes
+                if hasattr(pipe, 'enable_model_cpu_offload') and "sdxl" in model_id.lower():
                     pipe.enable_model_cpu_offload()
-                    print("✅ CPU offload habilitado")
+                    print("✅ CPU offload habilitado (modelo grande)")
                 
                 if hasattr(pipe, 'enable_vae_slicing'):
                     pipe.enable_vae_slicing()
                     print("✅ VAE slicing habilitado")
                 
+                # XFormers solo si está disponible y el modelo lo soporta
                 if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
                     try:
                         pipe.enable_xformers_memory_efficient_attention()
                         print("✅ XFormers memory efficient attention habilitado")
-                    except:
-                        print("⚠️ XFormers no disponible, usando atención estándar")
+                    except Exception as e:
+                        print(f"⚠️ XFormers no disponible: {e}")
+                        print("🔄 Usando atención estándar")
             
             current_model_id = model_id
             print(f"✅ Modelo {model_id} cargado exitosamente")
-            print(f"💾 Memoria utilizada: {torch.cuda.memory_allocated() / 1024**3:.2f} GB" if torch.cuda.is_available() else "💾 Memoria CPU")
+            
+            if torch.cuda.is_available():
+                memory_used = torch.cuda.memory_allocated() / 1024**3
+                memory_reserved = torch.cuda.memory_reserved() / 1024**3
+                print(f"💾 Memoria GPU utilizada: {memory_used:.2f} GB")
+                print(f"💾 Memoria GPU reservada: {memory_reserved:.2f} GB")
+                
+                # Verificar si la memoria es sospechosamente baja
+                if memory_used < 0.1:
+                    print("⚠️ ADVERTENCIA: Memoria GPU muy baja - posible problema de carga")
+            else:
+                print("💾 Memoria CPU")
             
         except Exception as e:
             print(f"❌ Error cargando modelo {model_id}: {e}")
             print(f"🔍 Tipo de error: {type(e).__name__}")
-            raise e
+            
+            # Intentar cargar sin variant fp16 si falló
+            if "variant" in str(e) and "fp16" in str(e):
+                print("🔄 Reintentando sin variant fp16...")
+                try:
+                    pipe = DiffusionPipeline.from_pretrained(
+                        model_id, 
+                        torch_dtype=torch_dtype,
+                        use_auth_token=HF_TOKEN if HF_TOKEN and ("flux" in model_id.lower() or "black-forest" in model_id.lower()) else None
+                    )
+                    pipe = pipe.to(device)
+                    current_model_id = model_id
+                    print(f"✅ Modelo {model_id} cargado exitosamente (sin fp16 variant)")
+                except Exception as e2:
+                    print(f"❌ Error en segundo intento: {e2}")
+                    raise e2
+            else:
+                raise e
     else:
         print(f"♻️ Modelo {model_id} ya está cargado, reutilizando...")
 
@@ -417,7 +461,8 @@ def infer(
             with torch.autocast(device_type='cuda', dtype=torch.float16):
                 print("⚡ Usando mixed precision para H200")
                 
-                image = pipe(
+                # Generar la imagen
+                result = pipe(
                     prompt=prompt,
                     negative_prompt=negative_prompt,
                     guidance_scale=final_guidance_scale,
@@ -426,10 +471,47 @@ def infer(
                     height=height,
                     generator=generator,
                     **additional_params
-                ).images[0]
+                )
+                
+                # Verificar que la imagen se generó correctamente
+                if hasattr(result, 'images') and len(result.images) > 0:
+                    image = result.images[0]
+                    
+                    # Verificar que la imagen no sea completamente negra
+                    if image is not None:
+                        # Convertir a numpy para verificar
+                        img_array = np.array(image)
+                        if img_array.size > 0:
+                            # Verificar si la imagen es completamente negra
+                            if np.all(img_array == 0) or np.all(img_array < 10):
+                                print("⚠️ ADVERTENCIA: Imagen generada es completamente negra")
+                                print("🔄 Reintentando con parámetros ajustados...")
+                                
+                                # Reintentar con parámetros más conservadores
+                                result = pipe(
+                                    prompt=prompt,
+                                    negative_prompt=negative_prompt,
+                                    guidance_scale=max(1.0, final_guidance_scale * 0.8),
+                                    num_inference_steps=max(10, final_inference_steps),
+                                    width=width,
+                                    height=height,
+                                    generator=generator
+                                )
+                                image = result.images[0]
+                            else:
+                                print("✅ Imagen generada correctamente")
+                        else:
+                            print("❌ Error: Imagen vacía")
+                            raise Exception("Imagen vacía generada")
+                    else:
+                        print("❌ Error: Imagen es None")
+                        raise Exception("Imagen es None")
+                else:
+                    print("❌ Error: No se generaron imágenes")
+                    raise Exception("No se generaron imágenes")
         else:
             # Fallback para CPU
-            image = pipe(
+            result = pipe(
                 prompt=prompt,
                 negative_prompt=negative_prompt,
                 guidance_scale=final_guidance_scale,
@@ -438,7 +520,8 @@ def infer(
                 height=height,
                 generator=generator,
                 **additional_params
-            ).images[0]
+            )
+            image = result.images[0]
 
         inference_time = time.time() - inference_start
         total_time = time.time() - start_time