Update app.py (#8)

- Update app.py (e0bb6b68b671f2e89c16d76c4772df4f913caed5)

app.py

@@ -7,7 +7,25 @@ from diffusers import DiffusionPipeline
 import io
 import base64
 from PIL import Image
-import json
+import logging
+from fastapi import FastAPI
+from pydantic import BaseModel
+
+# Configurar logging para depuração
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Inicializar FastAPI
+app = FastAPI()
+
+# Modelo para validação dos parâmetros da API
+class ImageRequest(BaseModel):
+    prompt: str
+    seed: int = 42
+    randomize_seed: bool = False
+    width: int = 1024
+    height: int = 1024
+    num_inference_steps: int = 4
 
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -19,6 +37,8 @@ MAX_IMAGE_SIZE = 2048
 
 @spaces.GPU()
 def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
+    logger.info(f"Chamando infer com prompt={prompt}, seed={seed}, randomize_seed={randomize_seed}, width={width}, height={height}, num_inference_steps={num_inference_steps}")
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
@@ -38,13 +58,21 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
     image.save(buffered, format="PNG")
     img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
     
-    # Retornar JSON com Base64 e seed
     return {"image_base64": f"data:image/png;base64,{img_str}", "seed": seed}
 
-# Função para a API personalizada
-def api_infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4):
-    result = infer(prompt, seed, randomize_seed, width, height, num_inference_steps)
-    return result  # Retorna diretamente o JSON
+# Endpoint FastAPI
+@app.post("/api/infer")
+async def api_infer(request: ImageRequest):
+    logger.info(f"Requisição API recebida: {request}")
+    result = infer(
+        prompt=request.prompt,
+        seed=request.seed,
+        randomize_seed=request.randomize_seed,
+        width=request.width,
+        height=request.height,
+        num_inference_steps=request.num_inference_steps
+    )
+    return result
 
 examples = [
     "a tiny astronaut hatching from an egg on the moon",
@@ -125,7 +153,6 @@ with gr.Blocks(css=css) as demo:
         output = infer(prompt, seed, randomize_seed, width, height, num_inference_steps)
         return output["image_base64"], output["seed"]
 
-    # Interface Gradio
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=format_output,
@@ -133,7 +160,5 @@ with gr.Blocks(css=css) as demo:
         outputs=[result, seed_output]
     )
 
-    # Endpoint personalizado para a API
-    demo.queue(api_name="infer_api").launch()
-
+# Iniciar o Gradio (sem queue, pois usamos FastAPI para a API)
 demo.launch()