Update app.py

app.py

@@ -1,17 +1,56 @@
-import gradio as gr
+from fastapi import FastAPI
+from pydantic import BaseModel
+from pyngrok import ngrok
+import nest_asyncio
+import uvicorn
+from diffusers import StableDiffusionPipeline
 import torch
-from diffusers import DiffusionPipeline
-from datasets import load_dataset
 from PIL import Image
+import io
+import base64
+
+# Instanciar a aplicação FastAPI
+app = FastAPI()
+
+# Carregar o modelo da Hugging Face
+model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"  # Substitua pelo ID do seu modelo
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+pipe = pipe.to("cuda")  # Use GPU
+
+# Definir o modelo de dados de entrada
+class InputData(BaseModel):
+    prompt: str
+
+# Endpoint inicial para verificar o funcionamento
+@app.get("/")
+def read_root():
+    return {"message": "API is running. Use /generate to create an image."}
+
+# Endpoint para gerar imagens
+@app.post("/generate")
+def generate_image(data: InputData):
+    # Gerar imagem com o modelo Hugging Face
+    image = pipe(data.prompt).images[0]
+    
+    # Converter a imagem para Base64
+    buffered = io.BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    
+    # Retornar o resultado
+    return {"result": f"Generated image for prompt: {data.prompt}", "image": img_str}
+
+# Expor o servidor publicamente com ngrok
+public_url = ngrok.connect(8000)
+print(f"Public URL: {public_url}")
+
+# Corrigir o loop de execução no Colab
+nest_asyncio.apply()
 
+# Iniciar o servidor
+uvicorn.run(app, host="0.0.0.0", port=8000)
 
-# Configurações do dispositivo para uso apenas da CPU
-device = "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Continuando com o modelo especificado
 
-# Carregar o pipeline configurado para CPU
-pipe = DiffusionPipeline.from_pretrained(model_repo_id)
-pipe = pipe.to(device)
 
 # Carregando o dataset do Hugging Face
 dataset = load_dataset("LEIDIA/Data_Womleimg")