yield current iterations

app.py

@@ -1,9 +1,15 @@
+import torch
 import gradio as gr
-from main import main
+from main import setup, execute_task
 from arguments import parse_args
 import os
 import shutil
 import glob
+import time
+import threading
+import argparse
+
+
 
 def list_iter_images(save_dir):
     # Specify the image extensions you want to search for
@@ -43,7 +49,16 @@ def clean_dir(save_dir):
     else:
         print(f"{save_dir} does not exist.")
 
-def generate_image(prompt, model, num_iterations, learning_rate, progress=gr.Progress(track_tqdm=True)):
+def start_over(gallery_state):
+    if gallery_state is not None:
+        gallery_state = None
+    return gallery_state, None, None, gr.update(visible=False)
+
+def setup_model(prompt, model, num_iterations, learning_rate, progress=gr.Progress(track_tqdm=True)):
+
+    """Clear CUDA memory before starting the training."""
+    torch.cuda.empty_cache()  # Free up cached memory
+    
     # Set up arguments
     args = parse_args()
     args.task = "single"
@@ -55,45 +70,83 @@ def generate_image(prompt, model, num_iterations, learning_rate, progress=gr.Pro
     args.save_dir = "./outputs"
     args.save_all_images = True
 
-    settings = (
-        f"{args.model}{'_' + args.prompt if args.task == 't2i-compbench' else ''}"
-        f"{'_no-optim' if args.no_optim else ''}_{args.seed if args.task != 'geneval' else ''}"
-        f"_lr{args.lr}_gc{args.grad_clip}_iter{args.n_iters}"
-        f"_reg{args.reg_weight if args.enable_reg else '0'}"
-        f"{'_pickscore' + str(args.pickscore_weighting) if args.enable_pickscore else ''}"
-        f"{'_clip' + str(args.clip_weighting) if args.enable_clip else ''}"
-        f"{'_hps' + str(args.hps_weighting) if args.enable_hps else ''}"
-        f"{'_imagereward' + str(args.imagereward_weighting) if args.enable_imagereward else ''}"
-        f"{'_aesthetic' + str(args.aesthetic_weighting) if args.enable_aesthetic else ''}"
-    )
+    args, trainer, device, dtype, shape, enable_grad, settings = setup(args)
+    loaded_setup = [args, trainer, device, dtype, shape, enable_grad, settings]
+
+    return None, loaded_setup
+
+def generate_image(setup_args, num_iterations):
+
+    args = setup_args[0]
+    trainer = setup_args[1]
+    device = setup_args[2]
+    dtype = setup_args[3]
+    shape = setup_args[4]
+    enable_grad = setup_args[5]
+
+    settings = setup_args[6]
 
     save_dir = f"{args.save_dir}/{args.task}/{settings}/{args.prompt}"
     clean_dir(save_dir)
     
     try:
-        # Run the main function with progress tracking
+        steps_completed = []
+        result_container = {"best_image": None, "total_init_rewards": None, "total_best_rewards": None}
+
+        # Define progress_callback that updates steps_completed
         def progress_callback(step):
-            progress(step / num_iterations, f"Iteration {step}/{num_iterations}")         
+            steps_completed.append(step)
 
-        best_image, total_init_rewards, total_best_rewards = main(args, progress_callback)
-        
-        # Return the path to the generated image
-        image_path = f"{save_dir}/best_image.png"
+        # Function to run main in a separate thread
+        def run_main():
+            result_container["best_image"], result_container["total_init_rewards"], result_container["total_best_rewards"] = execute_task(args, trainer, device, dtype, shape, enable_grad, settings, progress_callback)
         
-        if os.path.exists(image_path):
+        # Start main in a separate thread
+        main_thread = threading.Thread(target=run_main)
+        main_thread.start()
+
+        last_step_yielded = 0
+        while main_thread.is_alive() or last_step_yielded < num_iterations:
+            # Check if new steps have been completed
+            if steps_completed and steps_completed[-1] > last_step_yielded:
+                last_step_yielded = steps_completed[-1]
+                png_number = last_step_yielded - 1
+                # Get the image for this step
+                image_path = os.path.join(save_dir, f"{png_number}.png")
+                if os.path.exists(image_path):
+                    yield (image_path, f"Iteration {last_step_yielded}/{num_iterations} - Image saved", None)
+                else:
+                    yield (None, f"Iteration {last_step_yielded}/{num_iterations} - Image not found", None)
+            else:
+                # Small sleep to prevent busy waiting
+                time.sleep(0.1)
+
+        main_thread.join()
+
+        # After main is complete, yield the final image
+        final_image_path = os.path.join(save_dir, "best_image.png")
+        if os.path.exists(final_image_path):
             iter_images = list_iter_images(save_dir)
-            return image_path, f"Image generated successfully and saved at {image_path}", iter_images
+            yield (final_image_path, f"Final image saved at {final_image_path}", iter_images)
         else:
-            return None, "Image generation completed, but the file was not found.", None
-    
+            yield (None, "Image generation completed, but no final image was found.", None)
+
     except Exception as e:
-        return None, f"An error occurred: {str(e)}", None
+        yield (None, f"An error occurred: {str(e)}", None)
+
+def show_gallery_output(gallery_state):
+    if gallery_state is not None:
+        return gr.update(value=gallery_state, visible=True)
+    else:
+        return gr.update(value=None, visible=False)
 
 # Create Gradio interface
 title="# ReNO: Enhancing One-step Text-to-Image Models through Reward-based Noise Optimization"
 description="Enter a prompt to generate an image using ReNO. Adjust the model and parameters as needed."
 
 with gr.Blocks() as demo:
+    loaded_model_setup = gr.State()
+    gallery_state = gr.State()
     with gr.Column():
         gr.Markdown(title)
         gr.Markdown(description)
@@ -111,7 +164,9 @@ with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
                 prompt = gr.Textbox(label="Prompt")
-                chosen_model = gr.Dropdown(["sd-turbo", "sdxl-turbo", "pixart", "hyper-sd"], label="Model", value="sd-turbo")
+                with gr.Row():
+                    chosen_model = gr.Dropdown(["sd-turbo", "sdxl-turbo", "pixart", "hyper-sd"], label="Model", value="sd-turbo")
+                    model_status = gr.Textbox(label="model status", visible=False)
                 
                 with gr.Row():
                     n_iter = gr.Slider(minimum=10, maximum=100, step=10, value=50, label="Number of Iterations")
@@ -134,12 +189,24 @@ with gr.Blocks() as demo:
             with gr.Column():
                 output_image = gr.Image(type="filepath", label="Best Generated Image")
                 status = gr.Textbox(label="Status")
-                iter_gallery = gr.Gallery(label="Iterations", columns=4)
+                iter_gallery = gr.Gallery(label="Iterations", columns=4, visible=False)
 
     submit_btn.click(
-        fn = generate_image,
+        fn = start_over,
+        inputs =[gallery_state],
+        outputs = [gallery_state, output_image, status, iter_gallery]
+    ).then(
+        fn = setup_model,
         inputs = [prompt, chosen_model, n_iter, learning_rate],
-        outputs = [output_image, status, iter_gallery]
+        outputs =  [output_image, loaded_model_setup]
+    ).then(
+        fn = generate_image,
+        inputs = [loaded_model_setup, n_iter],
+        outputs = [output_image, status, gallery_state]
+    ).then(
+        fn = show_gallery_output,
+        inputs = [gallery_state],
+        outputs = iter_gallery
     )
 
 # Launch the app