otr

app.py

@@ -1,18 +1,43 @@
 import gradio as gr
-import requests
-import io
-from PIL import Image
+from diffusers import DiffusionPipeline
+import torch
 
-API_URL = "https://api-inference.huggingface.co/models/ZB-Tech/Text-to-Image"
-headers = {"Authorization": "Bearer HF_API_KEY"}
+# Cargar ambos modelos base y refiner
+base = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
+)
+base.to("cuda")
+refiner = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    text_encoder_2=base.text_encoder_2,
+    vae=base.vae,
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    variant="fp16",
+)
+refiner.to("cuda")
+
+# Definir cuántos pasos y qué porcentaje de pasos ejecutar en cada experto (80/20) aquí
+n_steps = 40
+high_noise_frac = 0.8
 
+# Definir función para generar imagen
 def generate_image(prompt):
-    payload = {"inputs": prompt}
-    response = requests.post(API_URL, headers=headers, json=payload)
-    image_bytes = response.content
-    image = Image.open(io.BytesIO(image_bytes))
+    image = base(
+        prompt=prompt,
+        num_inference_steps=n_steps,
+        denoising_end=high_noise_frac,
+        output_type="latent",
+    ).images
+    image = refiner(
+        prompt=prompt,
+        num_inference_steps=n_steps,
+        denoising_start=high_noise_frac,
+        image=image,
+    ).images[0]
     return image
 
+# Crear la interfaz de Gradio
 iface = gr.Interface(
     fn=generate_image,
     inputs=gr.Textbox(lines=5, label="Descripción de la imagen", placeholder="Introduce el texto aquí..."),
@@ -35,4 +60,5 @@ iface = gr.Interface(
     """
 )
 
+# Ejecutar la interfaz
 iface.launch()