Update app.py

app.py

@@ -1,40 +1,32 @@
 import torch
+from diffusers import AutoPipelineForText2Image
 import gradio as gr
-from diffusers import StableDiffusionPipeline
 
-# 1. Cargar el modelo
-# Usamos un modelo estándar de generación de imágenes
-model_id = "runwayml/stable-diffusion-v1-5"
+# Cargar modelo (se ejecuta una sola vez)
+pipe = AutoPipelineForText2Image.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    torch_dtype=torch.bfloat16
+)
+
+pipe = pipe.to("cuda")
 
-# Es fundamental usar torch_dtype=torch.float16 si tienes GPU para que sea más rápido.
-# Si estás en un entorno gratuito de Hugging Face sin GPU, el modelo usará la CPU automáticamente (aunque será más lento).
-try:
-    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
-    pipe = pipe.to("cpu") # Cambia "cpu" a "cuda" si compraste un entorno con GPU en Hugging Face
-except Exception as e:
-    pipe = None
-    print(f"Error al cargar el modelo: {e}")
+# Cargar LoRA
+pipe.load_lora_weights(
+    "enhanceaiteam/Flux-uncensored-v2",
+    weight_name="lora.safetensors"
+)
 
-# 2. Función de generación
-def generar_imagen(prompt):
-    if pipe is None:
-        return None # O podrías devolver una imagen de error por defecto
-    
-    # El pipeline devuelve una clase con una lista de imágenes. Tomamos la primera.
-    imagen_generada = pipe(prompt).images[0]
-    return imagen_generada
+# Función para generar imagen
+def generate(prompt):
+    image = pipe(prompt).images[0]
+    return image
 
-# 3. Interfaz de Gradio
+# Interfaz
 demo = gr.Interface(
-    fn=generar_imagen,
-    inputs=gr.Textbox(
-        label="¿Qué quieres dibujar?", 
-        placeholder="Ej: Un astronauta montando un caballo en Marte, estilo pintura al óleo"
-    ),
-    outputs=gr.Image(label="Imagen Generada", type="pil"),
-    title="Generador de Imágenes IA",
-    description="Escribe una descripción y la IA generará una imagen basándose en tu texto."
+    fn=generate,
+    inputs=gr.Textbox(label="Prompt"),
+    outputs=gr.Image(label="Resultado"),
+    title="FLUX + LoRA Demo"
 )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()