added gl ref

app.py

@@ -1,6 +1,13 @@
 import argparse
 import os
-
+import gc
+import psutil
+import threading
+from pathlib import Path
+import shutil
+import time
+import glob
+from datetime import datetime
 os.environ['CUDA_HOME'] = '/usr/local/cuda'
 os.environ['PATH'] = os.environ['PATH'] + ':/usr/local/cuda/bin'
 from datetime import datetime
@@ -117,7 +124,8 @@ def image_grid(imgs, rows, cols):
 
 
 args = parse_args()
-
+OUTPUT_DIR = "generated_images"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 # Mask-based CatVTON
 catvton_repo = "zhengchong/CatVTON"
 repo_path = snapshot_download(repo_id=catvton_repo)
@@ -148,6 +156,13 @@ pipeline_flux.load_lora_weights(
 )
 pipeline_flux.to("cuda", init_weight_dtype(args.mixed_precision))
 
+def save_generated_image(image, frame_no):
+    """Save generated image with timestamp and model name"""
+    filename = f"{frame_no}_frame.png"
+    filepath = os.path.join(OUTPUT_DIR, filename)
+    image.save(filepath)
+    return filepath
+    
 def print_image_info(img):
     # Basic attributes
     info = {
@@ -216,7 +231,7 @@ def process_video_frames(
     
     processed_frames = []
     print(f"processed_frames {len(frames)}")
-    for person_image in frames:
+    for index, person_image in enumerate(frames):
         result_image = proc_function_vidfl(
             person_image,
             cloth_image,
@@ -227,6 +242,7 @@ def process_video_frames(
         show_type
         )
         print_image_info(result_image)
+        save_generated_image(result_image,index)
         yield result_image
         processed_frames.append(result_image)
     
@@ -360,7 +376,28 @@ def submit_function_flux(
 def person_example_fn(image_path):
     return image_path
 
-
+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 update_gallery():
+        """Update the file gallery"""
+        files = get_generated_images()
+        return [
+            (f["path"], f"{f['name']}\n{f['date']}")
+            for f in files
+        ]
+    
 HEADER = """
 <h1 style="text-align: center;"> 🐈 CatVTON: Concatenation Is All You Need for Virtual Try-On with Diffusion Models </h1>
 
@@ -504,12 +541,19 @@ def app_gradio():
                         root_path = "resource/demo/example"
                         with gr.Column():
                             gal_output = gr.Gallery(label="Processed Frames")
+                            refresh_button = gr.Button("Refresh Gallery")
                         
 
                 image_path_vidflux.change(
                     person_example_fn, inputs=image_path_vidflux, outputs=person_image_vidflux
                 )
 
+                refresh_button.click(
+        fn=update_gallery,
+        inputs=[],
+        outputs=[file_gallery],
+    )
+
                 submit_flux.click(
                     process_video_frames,
                     [person_image_vidflux, cloth_image_vidflux, cloth_type, num_inference_steps_vidflux, guidance_scale_vidflux,