ensure no other thread is running

app.py

@@ -102,77 +102,77 @@ def generate_image(setup_args, num_iterations):
         steps_completed = []
         result_container = {"best_image": None, "total_init_rewards": None, "total_best_rewards": None}
         error_status = {"error_occurred": False}  # Shared dictionary to track error status
+        thread_status = {"running": False}  # Track whether a thread is already running
         
-        # Define progress_callback that updates steps_completed
         def progress_callback(step):
             # Limit redundant prints by checking the step number
             if not steps_completed or step > steps_completed[-1]:
                 steps_completed.append(step)
                 print(f"Progress: Step {step} completed.")
-
-        # Function to run main in a separate thread
+        
         def run_main():
+            thread_status["running"] = True  # Mark thread as running
             try:
-                # Call main and handle any potential OOM errors
-                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)
+                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)
             except torch.cuda.OutOfMemoryError as e:
-                # Handle CUDA OOM error
-                print("CUDA Out of Memory Error: ", e)
-                error_status["error_occurred"] = True  # Update status on error
+                print(f"CUDA Out of Memory Error: {e}")
+                error_status["error_occurred"] = True
             except RuntimeError as e:
                 if 'out of memory' in str(e):
-                    error_status["error_occurred"] = True  # Update status on error
+                    print(f"Runtime Error: {e}")
+                    error_status["error_occurred"] = True
                 else:
-                    raise  # Reraise if it's not a CUDA OOM error
+                    raise
+            finally:
+                thread_status["running"] = False  # Mark thread as completed
         
-        # 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() and not error_status["error_occurred"]:
-            # 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)
+        if not thread_status["running"]:  # Ensure no other thread is running
+            main_thread = threading.Thread(target=run_main)
+            main_thread.start()
+
+            last_step_yielded = 0
+            while main_thread.is_alive() and not error_status["error_occurred"]:
+                # 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:
-                    yield (None, f"Iteration {last_step_yielded}/{num_iterations} - Image not found", None)
-            else:
-                # Small sleep to prevent busy waiting
-                time.sleep(0.1)
+                    time.sleep(0.1)  # Sleep to prevent busy waiting
 
-        # If an error occurred, clean up resources and stop
-        if error_status["error_occurred"]:
-            torch.cuda.empty_cache()  # Free up cached memory
-            yield (None, "CUDA out of memory. Please reduce your batch size or image resolution.", None)
-        else:
-            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)
+            if error_status["error_occurred"]:
                 torch.cuda.empty_cache()  # Free up cached memory
-                yield (final_image_path, f"Final image saved at {final_image_path}", iter_images)
+                yield (None, "CUDA out of memory. Please reduce your batch size or image resolution.", None)
             else:
-                torch.cuda.empty_cache()  # Free up cached memory
-                yield (None, "Image generation completed, but no final image was found.", None)
+                main_thread.join()  # Ensure thread completion
+                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)
+                    torch.cuda.empty_cache()  # Free up cached memory
+                    yield (final_image_path, f"Final image saved at {final_image_path}", iter_images)
+                else:
+                    torch.cuda.empty_cache()  # Free up cached memory
+                    yield (None, "Image generation completed, but no final image was found.", None)
 
         torch.cuda.empty_cache()  # Free up cached memory
 
     except torch.cuda.OutOfMemoryError as e:
-        # Handle CUDA OOM error globally
+        print(f"Global CUDA Out of Memory Error: {e}")
         yield (None, "CUDA out of memory.", None)
     except RuntimeError as e:
         if 'out of memory' in str(e):
+            print(f"Runtime Error: {e}")
             yield (None, "CUDA out of memory.", None)
         else:
             yield (None, f"An error occurred: {str(e)}", None)
     except Exception as e:
+        print(f"Unexpected Error: {e}")
         yield (None, f"An unexpected error occurred: {str(e)}", None)
 
 def show_gallery_output(gallery_state):