Update app.py

app.py

@@ -7,6 +7,7 @@ from diffusers import DiffusionPipeline
 import io
 import base64
 from PIL import Image
+import json
 
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -40,6 +41,11 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
     # 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
+
 examples = [
     "a tiny astronaut hatching from an egg on the moon",
     "a cat holding a sign that says hello world",
@@ -119,6 +125,7 @@ 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,
@@ -126,4 +133,7 @@ with gr.Blocks(css=css) as demo:
         outputs=[result, seed_output]
     )
 
+    # Endpoint personalizado para a API
+    demo.queue(api_name="infer_api").launch()
+
 demo.launch()