new flow

app.py

@@ -15,12 +15,17 @@ import threading
 from pathlib import Path
 import shutil
 import time
+import glob
+from datetime import datetime
+from PIL import Image
 
 # Constants
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 TORCH_DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+OUTPUT_DIR = "generated_images"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 # Model configurations
 MODEL_CONFIGS = {
@@ -135,7 +140,30 @@ def load_pipeline(model_name):
     
     return pipe
 
-#@spaces.GPU(duration=180)
+def save_generated_image(image, model_name, prompt):
+    """Save generated image with timestamp and model name"""
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    # Create sanitized filename from prompt (first 30 chars)
+    prompt_part = "".join(c for c in prompt[:30] if c.isalnum() or c in (' ', '-', '_')).strip()
+    filename = f"{timestamp}_{model_name}_{prompt_part}.png"
+    filepath = os.path.join(OUTPUT_DIR, filename)
+    image.save(filepath)
+    return filepath
+
+def get_generated_images():
+    """Get list of generated images with their details"""
+    files = glob.glob(os.path.join(OUTPUT_DIR, "*.png"))
+    files.sort(key=os.path.getctime, reverse=True)  # Sort by creation time
+    return [
+        {
+            "path": f,
+            "name": os.path.basename(f),
+            "date": datetime.fromtimestamp(os.path.getctime(f)).strftime("%Y-%m-%d %H:%M:%S"),
+            "size": f"{os.path.getsize(f) / 1024:.1f} KB"
+        }
+        for f in files
+    ]
+
 def generate_image(
     model_name,
     prompt,
@@ -152,7 +180,6 @@ def generate_image(
         try:
             progress(0, desc=f"Loading {model_name} model...")
             
-            # Load model if not already loaded
             if model_name not in pipes:
                 pipes[model_name] = load_pipeline(model_name)
             
@@ -165,7 +192,6 @@ def generate_image(
             print(f"Generating image with {model_name}...")
             progress(0.3, desc=f"Generating image with {model_name}...")
             
-            # Generate image
             image = pipe(
                 prompt=prompt,
                 negative_prompt=negative_prompt,
@@ -176,9 +202,10 @@ def generate_image(
                 generator=generator,
             ).images[0]
             
-            progress(0.9, desc=f"Cleaning up {model_name} resources...")
+            filepath = save_generated_image(image, model_name, prompt)
+            print(f"Saved image to: {filepath}")
             
-            # Cleanup after generation
+            progress(0.9, desc=f"Cleaning up {model_name} resources...")
             deep_cleanup(model_name, pipe)
             
             progress(1.0, desc=f"Generation complete with {model_name}")
@@ -186,7 +213,6 @@ def generate_image(
             
         except Exception as e:
             print(f"Error with {model_name}: {str(e)}")
-            # Ensure cleanup happens even if generation fails
             if model_name in pipes:
                 deep_cleanup(model_name, pipes[model_name])
             raise e
@@ -198,7 +224,6 @@ css = """
     max-width: 1024px;
 }
 """
-#run_test_safe.zerogpu = True
 
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
@@ -263,17 +288,28 @@ with gr.Blocks(css=css) as demo:
                     value=40,
                 )
         
-        # Memory usage indicator
         memory_indicator = gr.Markdown("Current memory usage: 0 GB")
         
-        # Create tabs for each model
-        with gr.Tabs() as tabs:
-            results = {}
-            seeds = {}
-            for model_name in MODEL_CONFIGS.keys():
-                with gr.Tab(model_name):
-                    results[model_name] = gr.Image(label=f"{model_name} Result")
-                    seeds[model_name] = gr.Number(label="Seed used", visible=False)
+        with gr.Row():
+            with gr.Column(scale=2):
+                with gr.Tabs() as tabs:
+                    results = {}
+                    seeds = {}
+                    for model_name in MODEL_CONFIGS.keys():
+                        with gr.Tab(model_name):
+                            results[model_name] = gr.Image(label=f"{model_name} Result")
+                            seeds[model_name] = gr.Number(label="Seed used", visible=False)
+            
+            with gr.Column(scale=1):
+                gr.Markdown("### Generated Images")
+                file_gallery = gr.Gallery(
+                    label="Generated Images",
+                    show_label=False,
+                    elem_id="file_gallery",
+                    columns=2,
+                    height=400
+                )
+                refresh_button = gr.Button("Refresh Gallery")
         
         examples = [
             "A capybara wearing a suit holding a sign that reads Hello World",
@@ -281,15 +317,14 @@ with gr.Blocks(css=css) as demo:
         ]
         gr.Examples(examples=examples, inputs=[prompt])
 
-    def update_memory_usage():
-        """Update memory usage display"""
-        memory_gb = get_process_memory()
-        if torch.cuda.is_available():
-            cuda_memory_gb = torch.cuda.memory_allocated() / 1024 / 1024 / 1024
-            return f"Current memory usage: System RAM: {memory_gb:.2f} GB, CUDA: {cuda_memory_gb:.2f} GB"
-        return f"Current memory usage: System RAM: {memory_gb:.2f} GB"
+    def update_gallery():
+        """Update the file gallery"""
+        files = get_generated_images()
+        return [
+            (f["path"], f"{f['name']}\n{f['date']}")
+            for f in files
+        ]
 
-    # Handle generation for each model
     @spaces.GPU(duration=600)
     def generate_all(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress()):
         outputs = []
@@ -304,16 +339,16 @@ with gr.Blocks(css=css) as demo:
                 print(f"IMAGE GENERATED {model_name} {update_memory_usage()}")
                 outputs.extend([image, used_seed])
                 
-                # Update memory usage after each model
-                #memory_indicator.update(update_memory_usage())
-                
             except Exception as e:
                 outputs.extend([None, None])
                 print(f"Error generating with {model_name}: {str(e)}")
         
+        # Update the gallery after generation
+        gallery_images = update_gallery()
+        file_gallery.update(value=gallery_images)
+        
         return outputs
 
-    # Set up the generation trigger
     output_components = []
     for model_name in MODEL_CONFIGS.keys():
         output_components.extend([results[model_name], seeds[model_name]])
@@ -333,5 +368,17 @@ with gr.Blocks(css=css) as demo:
         outputs=output_components,
     )
 
+    refresh_button.click(
+        fn=update_gallery,
+        inputs=[],
+        outputs=[file_gallery],
+    )
+
+    demo.load(
+        fn=update_gallery,
+        inputs=[],
+        outputs=[file_gallery],
+    )
+
 if __name__ == "__main__":
     demo.launch()