Update app.py

app.py

@@ -1,13 +1,20 @@
 import torch
 from diffusers import FluxPipeline
 
-# Cargar el pipeline con cuantización de 4 bits
+# 1. Definir el dispositivo
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# 2. Cargar el modelo con cuantización (si tienes bitsandbytes instalado)
+# Si sigues teniendo error de memoria, intenta usar "black-forest-labs/FLUX.1-schnell"
 pipe = FluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev", 
     torch_dtype=torch.bfloat16,
-    load_in_4bit=True, # Esta es la clave para reducir la memoria
-    device_map="auto"
+    # load_in_4bit=True requiere configuraciones adicionales en diffusers,
+    # prueba primero sin él, usando solamente bfloat16 para ahorrar VRAM.
 )
 
-# Luego carga tu LoRA como ya sabías
+# 3. Mover el pipeline a la GPU después de cargarlo
+pipe.to(device)
+
+# 4. Cargar el LoRA
 pipe.load_lora_weights("joyfox/Kontext-Cosplay-Lora", weight_name="kontext_cosplay_lora_20000.safetensors")