Update app.py

app.py

@@ -3,6 +3,8 @@ import numpy as np
 import random
 import spaces
 import torch
+import os
+import time
 from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
 from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
 from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
@@ -40,7 +42,22 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidan
         ):
             final_img = img
     
-    return final_img, seed
+    # Save image to /tmp directory and get the path
+    if final_img is not None:
+        # Create a unique filename with timestamp
+        timestamp = int(time.time())
+        filename = f"flux_generated_{timestamp}_{seed}.png"
+        image_path = os.path.join("/tmp", filename)
+        
+        # Ensure /tmp directory exists (it should by default on Linux)
+        os.makedirs("/tmp", exist_ok=True)
+        
+        # Save the image
+        final_img.save(image_path)
+        
+        return final_img, seed, image_path
+    
+    return final_img, seed, ""
     
 examples = [
     "a tiny astronaut hatching from an egg on the moon",
@@ -77,6 +94,14 @@ with gr.Blocks(css=css) as demo:
         
         result = gr.Image(label="Result", show_label=False)
         
+        # Add a text box to display the saved image path
+        image_path_output = gr.Textbox(
+            label="Saved Image Path",
+            placeholder="Generated image path will appear here...",
+            interactive=False,
+            show_copy_button=True
+        )
+        
         with gr.Accordion("Advanced Settings", open=False):
             
             seed = gr.Slider(
@@ -129,7 +154,7 @@ with gr.Blocks(css=css) as demo:
             examples = examples,
             fn = infer,
             inputs = [prompt],
-            outputs = [result, seed],
+            outputs = [result, seed, image_path_output],
             cache_examples="lazy"
         )
 
@@ -137,7 +162,7 @@ with gr.Blocks(css=css) as demo:
         triggers=[run_button.click, prompt.submit],
         fn = infer,
         inputs = [prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs = [result, seed]
+        outputs = [result, seed, image_path_output]
     )
 
 demo.launch()