Update app.py

app.py

@@ -4,9 +4,9 @@ import random
 import spaces
 import torch
 from diffusers import DiffusionPipeline
-import os
-import uuid
-from datetime import datetime
+import io
+import base64
+from PIL import Image
 
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -16,10 +16,6 @@ pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", tor
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-# Diretório para salvar imagens (use /tmp para Hugging Face Spaces)
-OUTPUT_DIR = "/tmp"
-os.makedirs(OUTPUT_DIR, exist_ok=True)
-
 @spaces.GPU()
 def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
@@ -36,18 +32,13 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
         guidance_scale=0.0
     ).images[0]
     
-    # Gerar nome de arquivo único
-    filename = f"image_{uuid.uuid4().hex}.png"
-    filepath = os.path.join(OUTPUT_DIR, filename)
-    
-    # Salvar a imagem
-    image.save(filepath)
-    
-    # Construir o URL público (ajuste conforme a URL do seu Space)
-    space_url = "https://eder0782-flux-image-generator.hf.space"
-    image_url = f"{space_url}/file/{filename}"
+    # Converter a imagem para Base64
+    buffered = io.BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
     
-    return {"image_url": image_url, "seed": seed}
+    # Retornar JSON com Base64 e seed
+    return {"image_base64": f"data:image/png;base64,{img_str}", "seed": seed}
 
 examples = [
     "a tiny astronaut hatching from an egg on the moon",
@@ -80,6 +71,7 @@ with gr.Blocks(css=css) as demo:
             run_button = gr.Button("Run", scale=0)
         
         result = gr.Image(label="Result", show_label=False)
+        seed_output = gr.Number(label="Seed", show_label=True)
         
         with gr.Accordion("Advanced Settings", open=False):
             seed = gr.Slider(
@@ -117,20 +109,21 @@ with gr.Blocks(css=css) as demo:
             examples=examples,
             fn=infer,
             inputs=[prompt],
-            outputs=[result, seed],
-            cache_examples="lazy"
+            outputs=[result, seed_output],
+            cache_examples=True,
+            cache_mode="lazy"
         )
 
-    # Ajustar a saída para a interface
-    def format_output(output):
-        return output["image_url"], output["seed"]
+    # Função para formatar a saída para a interface
+    def format_output(prompt, seed, randomize_seed, width, height, num_inference_steps):
+        output = infer(prompt, seed, randomize_seed, width, height, num_inference_steps)
+        return output["image_base64"], output["seed"]
 
     gr.on(
         triggers=[run_button.click, prompt.submit],
-        fn=infer,
+        fn=format_output,
         inputs=[prompt, seed, randomize_seed, width, height, num_inference_steps],
-        outputs=[result, seed],
-        _js=format_output
+        outputs=[result, seed_output]
     )
 
 demo.launch()