Update app.py

app.py

@@ -1,33 +1,36 @@
 import gradio as gr
-from safetensors.torch import load_file
-import os
-import requests
+from diffusers import StableDiffusionPipeline
+import torch
 
-# URLs y rutas de archivos
-MODEL_URL = "https://huggingface.co/datasets/uruguayai/fooocus/resolve/main/juggernautXL_v8Rundiffusion.safetensors"
-MODEL_PATH = "/home/user/app/juggernautXL_v8Rundiffusion.safetensors"
-
-def download_file(url, path):
-    if not os.path.exists(path):
-        response = requests.get(url)
-        if response.status_code == 200:
-            with open(path, 'wb') as f:
-                f.write(response.content)
-            print(f"Downloaded {path}")
-        else:
-            raise Exception(f"Failed to download {url}, status code: {response.status_code}")
+# Configura el modelo, puedes ajustar el modelo a tu necesidad específica.
+MODEL_NAME = "CompVis/stable-diffusion-v1-4"  # Puedes cambiarlo a tu modelo subido en HF
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def load_model():
-    download_file(MODEL_URL, MODEL_PATH)
-    model = load_file(MODEL_PATH)  # Cargar modelo con safetensors
-    return model
+    # Cargar el modelo de Stable Diffusion
+    pipe = StableDiffusionPipeline.from_pretrained(MODEL_NAME, torch_dtype=torch.float16).to(device)
+    return pipe
+
+# Cargar el modelo al iniciar la aplicación
+model = load_model()
 
-def predict(input_text):
-    # Función de predicción simple
-    return f"Model loaded and ready: {input_text}"
+def generate_image(prompt):
+    # Generar imagen usando el modelo cargado
+    if not prompt:
+        return None
+    with torch.no_grad():
+        image = model(prompt).images[0]  # Genera la imagen
+    return image
 
-iface = gr.Interface(fn=predict, inputs="text", outputs="text")
+# Configurar la interfaz de Gradio
+iface = gr.Interface(
+    fn=generate_image,
+    inputs=gr.Textbox(label="Enter your prompt"),
+    outputs=gr.Image(label="Generated Image"),
+    title="AI Image Generator",
+    description="Generate images from text prompts using Stable Diffusion."
+)
 
 if __name__ == "__main__":
-    load_model()  # Asegúrate de que el modelo esté cargado
+    # Lanzar la interfaz
     iface.launch(server_name="0.0.0.0", server_port=7860)