Update app.py

app.py

@@ -1,419 +1,409 @@
-import json
-import logging
-import os
-
-import blobfile as bf
 import torch
 import gc
-from datasets import load_dataset
-from pytorch_lightning import seed_everything
-from tqdm import tqdm
-
+import gradio as gr
+from main import setup, execute_task
 from arguments import parse_args
-from models import get_model, get_multi_apply_fn
-from rewards import get_reward_losses
-from training import LatentNoiseTrainer, get_optimizer
-
+import os
+import shutil
+import glob
+import time
+import threading
+import argparse
+
+def list_iter_images(save_dir):
+    # Specify only PNG images
+    image_extension = 'png'
+
+    # Create a list to store the image file paths
+    image_paths = []
+
+    # Use glob to find all PNG image files
+    all_images = glob.glob(os.path.join(save_dir, f'*.{image_extension}'))
+
+    # Filter out 'best_image.png'
+    image_paths = [img for img in all_images if os.path.basename(img) != 'best_image.png']
+
+    return image_paths
+
+def clean_dir(save_dir):
+    # Check if the directory exists
+    if os.path.exists(save_dir):
+        # Check if the directory contains any files
+        if len(os.listdir(save_dir)) > 0:
+            # If it contains files, delete all files in the directory
+            for filename in os.listdir(save_dir):
+                file_path = os.path.join(save_dir, filename)
+                try:
+                    if os.path.isfile(file_path) or os.path.islink(file_path):
+                        os.unlink(file_path)  # Remove file or symbolic link
+                    elif os.path.isdir(file_path):
+                        shutil.rmtree(file_path)  # Remove directory and its contents
+                except Exception as e:
+                    print(f"Failed to delete {file_path}. Reason: {e}")
+            print(f"All files in {save_dir} have been deleted.")
+        else:
+            print(f"{save_dir} exists but is empty.")
+    else:
+        print(f"{save_dir} does not exist.")
 
-import torch
-import gc
+def start_over(gallery_state):
+    torch.cuda.empty_cache()  # Free up cached memory
+    gc.collect()
+    if gallery_state is not None:
+        gallery_state = None
+    return gallery_state, None, None, gr.update(visible=False)
 
-def clear_gpu():
-    """Clear GPU memory by removing tensors, freeing cache, and moving data to CPU."""
-    # List memory usage before clearing
-    print(f"Memory allocated before clearing: {torch.cuda.memory_allocated() / (1024 ** 2)} MB")
-    print(f"Memory reserved before clearing: {torch.cuda.memory_reserved() / (1024 ** 2)} MB")
+def setup_model(loaded_model_setup, prompt, model, seed, num_iterations, enable_hps, hps_w, enable_imagereward, imgrw_w, enable_pickscore, pcks_w, enable_clip, clip_w, learning_rate, progress=gr.Progress(track_tqdm=True)):
+    gr.Info(f"Loading {model} model ...")
     
-    # Force the garbage collector to free unreferenced objects
-    gc.collect()
+    if prompt is None or prompt == "":
+        raise gr.Error("You forgot to provide a prompt !")
 
-    # Move any bound tensors back to CPU if needed
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()  # Free up the cached memory
-        torch.cuda.ipc_collect()  # Clear any cross-process memory
+    print(f"LOADED_MODEL SETUP: {loaded_model_setup}")
     
-    print(f"Memory allocated after clearing: {torch.cuda.memory_allocated() / (1024 ** 2)} MB")
-    print(f"Memory reserved after clearing: {torch.cuda.memory_reserved() / (1024 ** 2)} MB")
-
-def unload_previous_model_if_needed(loaded_model_setup):
-    """Unload the current model from the GPU and free resources if a new model is being loaded."""
-    if loaded_model_setup is not None:
-        print("Unloading previous model from GPU to free memory.")
-        previous_model = loaded_model_setup[7]  # Assuming pipe is at position [7] in the setup
-        if hasattr(previous_model, 'to') and loaded_model_setup[0].model != "flux":
-            previous_model.to('cpu')  # Move model to CPU to free GPU memory
-        del previous_model  # Delete the reference to the model
-        clear_gpu()  # Clear all remaining GPU memory
-
-def setup(args, loaded_model_setup=None):
-    seed_everything(args.seed)
-    bf.makedirs(f"{args.save_dir}/logs/{args.task}")
+    """Clear CUDA memory before starting the training."""
+    torch.cuda.empty_cache()  # Free up cached memory
+    gc.collect()
     
-    # Set up logging and name settings
-    logger = logging.getLogger()
-    logger.handlers.clear()  # Clear existing handlers
-    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 ''}"
-    )
+    # Set up arguments
+    args = parse_args()
+    args.task = "single"
+    args.prompt = prompt
+    args.model = model
+    args.seed = seed
+    args.n_iters = num_iterations
+    args.lr = learning_rate
+    args.cache_dir = "./HF_model_cache"
+    args.save_dir = "./outputs"
+    args.save_all_images = True
+
+    if enable_hps is True:
+        args.disable_hps = False
+        args.hps_weighting = hps_w
     
-    file_stream = open(f"{args.save_dir}/logs/{args.task}/{settings}.txt", "w")
-    handler = logging.StreamHandler(file_stream)
-    formatter = logging.Formatter("%(asctime)s - %(message)s")
-    handler.setFormatter(formatter)
-    logger.addHandler(handler)
-    logger.setLevel("INFO")
-    consoleHandler = logging.StreamHandler()
-    consoleHandler.setFormatter(formatter)
-    logger.addHandler(consoleHandler)
+    if enable_imagereward is True:
+        args.disable_imagereward = False
+        args.imagereward_weighting = imgrw_w
     
-    logging.info(args)
+    if enable_pickscore is True:
+        args.disable_pickscore = False
+        args.pickscore_weighting = pcks_w
     
-    if args.device_id is not None:
-        logging.info(f"Using CUDA device {args.device_id}")
-        os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-        os.environ["CUDA_VISIBLE_DEVICES"] = args.device_id
-
-    device = torch.device("cuda")
-    dtype = torch.float16 if args.dtype == "float16" else torch.float32
-
-    # If args.model is the same as the one in loaded_model_setup, reuse the trainer and pipe
-    if loaded_model_setup and args.model == loaded_model_setup[0].model:
-        print(f"Reusing model {args.model} from loaded setup.")
-        trainer = loaded_model_setup[1]  # Trainer is at position 1 in loaded_model_setup
-        
-        # Update trainer with the new arguments
-        trainer.n_iters = args.n_iters
-        trainer.n_inference_steps = args.n_inference_steps
-        trainer.seed = args.seed
-        trainer.save_all_images = args.save_all_images
-        trainer.no_optim = args.no_optim
-        trainer.regularize = args.enable_reg
-        trainer.regularization_weight = args.reg_weight
-        trainer.grad_clip = args.grad_clip
-        trainer.log_metrics = args.task == "single" or not args.no_optim
-        trainer.imageselect = args.imageselect
+    if enable_clip is True:
+        args.disable_clip = False
+        args.clip_weighting = clip_w
+
+    if model == "flux":
+        args.cpu_offloading = True
+        args.enable_multi_apply = True
+        args.multi_step_model = "flux"
+
+    # Check if args are the same as the loaded_model_setup except for the prompt
+    if loaded_model_setup and hasattr(loaded_model_setup[0], '__dict__'):
+        previous_args = loaded_model_setup[0]
         
-        # Get latents (this step is still required)
-        if args.model == "flux":
-            shape = (1, 16 * 64, 64)
-        elif args.model != "pixart":
-            height = trainer.model.unet.config.sample_size * trainer.model.vae_scale_factor
-            width = trainer.model.unet.config.sample_size * trainer.model.vae_scale_factor
-            shape = (
-                1,
-                trainer.model.unet.in_channels,
-                height // trainer.model.vae_scale_factor,
-                width // trainer.model.vae_scale_factor,
-            )
-        else:
-            height = trainer.model.transformer.config.sample_size * trainer.model.vae_scale_factor
-            width = trainer.model.transformer.config.sample_size * trainer.model.vae_scale_factor
-            shape = (
-                1,
-                trainer.model.transformer.config.in_channels,
-                height // trainer.model.vae_scale_factor,
-                width // trainer.model.vae_scale_factor,
-            )
+        # Exclude 'prompt' from comparison
+        new_args_dict = {k: v for k, v in args.__dict__.items() if k != 'prompt'}
+        prev_args_dict = {k: v for k, v in previous_args.__dict__.items() if k != 'prompt'}
         
-        pipe = loaded_model_setup[7]
-        enable_grad = not args.no_optim
-
-        return args, trainer, device, dtype, shape, enable_grad, settings, pipe
-
-    # Unload previous model and clear GPU resources
-    unload_previous_model_if_needed(loaded_model_setup)
-
-    # Proceed with full model loading if args.model is different
-    print(f"Loading new model: {args.model}")
-    
-    # Get reward losses
-    reward_losses = get_reward_losses(args, dtype, device, args.cache_dir)
-
-    # Get model and noise trainer
-    pipe = get_model(
-        args.model, dtype, device, args.cache_dir, args.memsave, args.cpu_offloading
-    )
-
-    # Final memory cleanup after model loading
-    torch.cuda.empty_cache()
-    gc.collect()
-
-    trainer = LatentNoiseTrainer(
-        reward_losses=reward_losses,
-        model=pipe,
-        n_iters=args.n_iters,
-        n_inference_steps=args.n_inference_steps,
-        seed=args.seed,
-        save_all_images=args.save_all_images,
-        device=device if not args.cpu_offloading else 'cpu',  # Use CPU if offloading is enabled
-        no_optim=args.no_optim,
-        regularize=args.enable_reg,
-        regularization_weight=args.reg_weight,
-        grad_clip=args.grad_clip,
-        log_metrics=args.task == "single" or not args.no_optim,
-        imageselect=args.imageselect,
-    )
-
-    # Create latents
-    if args.model == "flux":
-        shape = (1, 16 * 64, 64)
-    elif args.model != "pixart":
-        height = pipe.unet.config.sample_size * pipe.vae_scale_factor
-        width = pipe.unet.config.sample_size * pipe.vae_scale_factor
-        shape = (
-            1,
-            pipe.unet.in_channels,
-            height // pipe.vae_scale_factor,
-            width // pipe.vae_scale_factor,
-        )
-    else:
-        height = pipe.transformer.config.sample_size * pipe.vae_scale_factor
-        width = pipe.transformer.config.sample_size * pipe.vae_scale_factor
-        shape = (
-            1,
-            pipe.transformer.config.in_channels,
-            height // pipe.vae_scale_factor,
-            width // pipe.vae_scale_factor,
-        )
+        if new_args_dict == prev_args_dict:
+            # If the arguments (excluding prompt) are the same, reuse the loaded setup
+            print(f"Arguments (excluding prompt) are the same, reusing loaded setup for {model} model.")
+            
+            # Update the prompt in the loaded_model_setup
+            loaded_model_setup[0].prompt = prompt
+            
+            yield f"{model} model already loaded with the same configuration.", loaded_model_setup      
+
+    # Attempt to set up the model
+    try:
+        # If other args differ, proceed with the setup
+        args, trainer, device, dtype, shape, enable_grad, settings, pipe = setup(args, loaded_model_setup)
+        new_loaded_setup = [args, trainer, device, dtype, shape, enable_grad, settings, pipe]
+        yield f"{model} model loaded successfully!", new_loaded_setup
     
-    enable_grad = not args.no_optim
+    except Exception as e:
+        print(f"Failed to load {model} model: {e}.")
+        yield f"Failed to load {model} model: {e}. You can try again, as it usually finally loads on the second try :)", None
+       
 
-    # Final memory cleanup
+def generate_image(setup_args, num_iterations):
     torch.cuda.empty_cache()  # Free up cached memory
     gc.collect()
 
-    
-
-    return args, trainer, device, dtype, shape, enable_grad, settings, pipe
+    gr.Info(f"Executing iterations task ...")
 
+    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]
+    print(f"SETTINGS: {settings}")
 
+    pipe = setup_args[7]
 
-def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pipe, progress_callback=None):
+    save_dir = f"{args.save_dir}/{args.task}/{settings}/{args.prompt[:150]}"
+    clean_dir(save_dir)
     
-    if args.task == "single":
-        # Attempt to move the model to GPU if model is not Flux
-        if args.model != "flux":
-            if pipe.device != torch.device('cuda'):
-                pipe.to(device, dtype)
-        else:
-            print(f"PIPE:{pipe}")
-            
-            
-            if args.cpu_offloading:
-                pipe.enable_sequential_cpu_offload()
-
-            #if pipe.device != torch.device('cuda'):
-            #    pipe.to(device, dtype)
+    try:
+        torch.cuda.empty_cache()  # Free up cached memory
+        gc.collect()
+        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
         
-        if args.enable_multi_apply:
-            
-            multi_apply_fn = get_multi_apply_fn(
-                model_type=args.multi_step_model,
-                seed=args.seed,
-                pipe=pipe,
-                cache_dir=args.cache_dir,
-                device=device if not args.cpu_offloading else 'cpu',
-                dtype=dtype,
-            )
-        else:
-            multi_apply_fn = None    
-    
+        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.")
+        
+        def run_main():
+            thread_status["running"] = True  # Mark thread as running
+            try:
+                execute_task(
+                    args, trainer, device, dtype, shape, enable_grad, settings, pipe, progress_callback
+                )
+            except torch.cuda.OutOfMemoryError as e:
+                print(f"CUDA Out of Memory Error: {e}")
+                error_status["error_occurred"] = True
+            except RuntimeError as e:
+                if 'out of memory' in str(e):
+                    print(f"Runtime Error: {e}")
+                    error_status["error_occurred"] = True
+                else:
+                    raise
+            finally:
+                thread_status["running"] = False  # Mark thread as completed
+        
+        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:
+                    time.sleep(0.1)  # Sleep to prevent busy waiting
+
+            if error_status["error_occurred"]:
+                torch.cuda.empty_cache()  # Free up cached memory
+                gc.collect()
+                yield (None, "CUDA out of memory. Please reduce your batch size or image resolution.", None)
+            else:
+                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
+                    gc.collect()
+                    time.sleep(0.5)
+                    yield (final_image_path, f"Final image saved at {final_image_path}", iter_images)
+                else:
+                    torch.cuda.empty_cache()  # Free up cached memory
+                    gc.collect()
+                    yield (None, "Image generation completed, but no final image was found.", None)
+
         torch.cuda.empty_cache()  # Free up cached memory
         gc.collect()
-    
 
-        init_latents = torch.randn(shape, device=device, dtype=dtype)
-        latents = torch.nn.Parameter(init_latents, requires_grad=enable_grad)
-        optimizer = get_optimizer(args.optim, latents, args.lr, args.nesterov)
-        save_dir = f"{args.save_dir}/{args.task}/{settings}/{args.prompt[:150]}"
-        os.makedirs(f"{save_dir}", exist_ok=True)
-        init_image, best_image, total_init_rewards, total_best_rewards = trainer.train(
-            latents, args.prompt, optimizer, save_dir, multi_apply_fn, progress_callback=progress_callback
-        )
-        best_image.save(f"{save_dir}/best_image.png")
-        #init_image.save(f"{save_dir}/init_image.png")
-
-    elif args.task == "example-prompts":
-        fo = open("assets/example_prompts.txt", "r")
-        prompts = fo.readlines()
-        fo.close()
-        for i, prompt in tqdm(enumerate(prompts)):
-            # Get new latents and optimizer
-            init_latents = torch.randn(shape, device=device, dtype=dtype)
-            latents = torch.nn.Parameter(init_latents, requires_grad=enable_grad)
-            optimizer = get_optimizer(args.optim, latents, args.lr, args.nesterov)
-
-            prompt = prompt.strip()
-            name = f"{i:03d}_{prompt[:150]}.png"
-            save_dir = f"{args.save_dir}/{args.task}/{settings}/{name}"
-            os.makedirs(save_dir, exist_ok=True)
-            init_image, best_image, init_rewards, best_rewards = trainer.train(
-                latents, prompt, optimizer, save_dir, multi_apply_fn
-            )
-            if i == 0:
-                total_best_rewards = {k: 0.0 for k in best_rewards.keys()}
-                total_init_rewards = {k: 0.0 for k in best_rewards.keys()}
-            for k in best_rewards.keys():
-                total_best_rewards[k] += best_rewards[k]
-                total_init_rewards[k] += init_rewards[k]
-            best_image.save(f"{save_dir}/best_image.png")
-            init_image.save(f"{save_dir}/init_image.png")
-            logging.info(f"Initial rewards: {init_rewards}")
-            logging.info(f"Best rewards: {best_rewards}")
-        for k in total_best_rewards.keys():
-            total_best_rewards[k] /= len(prompts)
-            total_init_rewards[k] /= len(prompts)
-
-        # save results to directory
-        with open(f"{args.save_dir}/example-prompts/{settings}/results.txt", "w") as f:
-            f.write(
-                f"Mean initial all rewards: {total_init_rewards}\n"
-                f"Mean best all rewards: {total_best_rewards}\n"
-            )
-    elif args.task == "t2i-compbench":
-        prompt_list_file = f"../T2I-CompBench/examples/dataset/{args.prompt}.txt"
-        fo = open(prompt_list_file, "r")
-        prompts = fo.readlines()
-        fo.close()
-        os.makedirs(f"{args.save_dir}/{args.task}/{settings}/samples", exist_ok=True)
-        for i, prompt in tqdm(enumerate(prompts)):
-            # Get new latents and optimizer
-            init_latents = torch.randn(shape, device=device, dtype=dtype)
-            latents = torch.nn.Parameter(init_latents, requires_grad=enable_grad)
-            optimizer = get_optimizer(args.optim, latents, args.lr, args.nesterov)
-
-            prompt = prompt.strip()
-            init_image, best_image, init_rewards, best_rewards = trainer.train(
-                latents, prompt, optimizer, None, multi_apply_fn
-            )
-            if i == 0:
-                total_best_rewards = {k: 0.0 for k in best_rewards.keys()}
-                total_init_rewards = {k: 0.0 for k in best_rewards.keys()}
-            for k in best_rewards.keys():
-                total_best_rewards[k] += best_rewards[k]
-                total_init_rewards[k] += init_rewards[k]
-            name = f"{prompt}_{i:06d}.png"
-            best_image.save(f"{args.save_dir}/{args.task}/{settings}/samples/{name}")
-            logging.info(f"Initial rewards: {init_rewards}")
-            logging.info(f"Best rewards: {best_rewards}")
-        for k in total_best_rewards.keys():
-            total_best_rewards[k] /= len(prompts)
-            total_init_rewards[k] /= len(prompts)
-    elif args.task == "parti-prompts":
-        parti_dataset = load_dataset("nateraw/parti-prompts", split="train")
-        total_reward_diff = 0.0
-        total_best_reward = 0.0
-        total_init_reward = 0.0
-        total_improved_samples = 0
-        for index, sample in enumerate(parti_dataset):
-            os.makedirs(
-                f"{args.save_dir}/{args.task}/{settings}/{index}", exist_ok=True
-            )
-            prompt = sample["Prompt"]
-            init_image, best_image, init_rewards, best_rewards = trainer.train(
-                latents, prompt, optimizer, multi_apply_fn
-            )
-            best_image.save(
-                f"{args.save_dir}/{args.task}/{settings}/{index}/best_image.png"
-            )
-            open(
-                f"{args.save_dir}/{args.task}/{settings}/{index}/prompt.txt", "w"
-            ).write(
-                f"{prompt} \n Initial Rewards: {init_rewards} \n Best Rewards: {best_rewards}"
-            )
-            logging.info(f"Initial rewards: {init_rewards}")
-            logging.info(f"Best rewards: {best_rewards}")
-            initial_reward = init_rewards[args.benchmark_reward]
-            best_reward = best_rewards[args.benchmark_reward]
-            total_reward_diff += best_reward - initial_reward
-            total_best_reward += best_reward
-            total_init_reward += initial_reward
-            if best_reward < initial_reward:
-                total_improved_samples += 1
-            if i == 0:
-                total_best_rewards = {k: 0.0 for k in best_rewards.keys()}
-                total_init_rewards = {k: 0.0 for k in best_rewards.keys()}
-            for k in best_rewards.keys():
-                total_best_rewards[k] += best_rewards[k]
-                total_init_rewards[k] += init_rewards[k]
-            # Get new latents and optimizer
-            init_latents = torch.randn(shape, device=device, dtype=dtype)
-            latents = torch.nn.Parameter(init_latents, requires_grad=enable_grad)
-            optimizer = get_optimizer(args.optim, latents, args.lr, args.nesterov)
-        improvement_percentage = total_improved_samples / parti_dataset.num_rows
-        mean_best_reward = total_best_reward / parti_dataset.num_rows
-        mean_init_reward = total_init_reward / parti_dataset.num_rows
-        mean_reward_diff = total_reward_diff / parti_dataset.num_rows
-        logging.info(
-            f"Improvement percentage: {improvement_percentage:.4f}, "
-            f"mean initial reward: {mean_init_reward:.4f}, "
-            f"mean best reward: {mean_best_reward:.4f}, "
-            f"mean reward diff: {mean_reward_diff:.4f}"
-        )
-        for k in total_best_rewards.keys():
-            total_best_rewards[k] /= len(parti_dataset)
-            total_init_rewards[k] /= len(parti_dataset)
-        # save results
-        os.makedirs(f"{args.save_dir}/parti-prompts/{settings}", exist_ok=True)
-        with open(f"{args.save_dir}/parti-prompts/{settings}/results.txt", "w") as f:
-            f.write(
-                f"Mean improvement: {improvement_percentage:.4f}, "
-                f"mean initial reward: {mean_init_reward:.4f}, "
-                f"mean best reward: {mean_best_reward:.4f}, "
-                f"mean reward diff: {mean_reward_diff:.4f}\n"
-                f"Mean initial all rewards: {total_init_rewards}\n"
-                f"Mean best all rewards: {total_best_rewards}"
-            )
-    elif args.task == "geneval":
-        prompt_list_file = "../geneval/prompts/evaluation_metadata.jsonl"
-        with open(prompt_list_file) as fp:
-            metadatas = [json.loads(line) for line in fp]
-        outdir = f"{args.save_dir}/{args.task}/{settings}"
-        for index, metadata in enumerate(metadatas):
-            # Get new latents and optimizer
-            init_latents = torch.randn(shape, device=device, dtype=dtype)
-            latents = torch.nn.Parameter(init_latents, requires_grad=True)
-            optimizer = get_optimizer(args.optim, latents, args.lr, args.nesterov)
-
-            prompt = metadata["prompt"]
-            init_image, best_image, init_rewards, best_rewards = trainer.train(
-                latents, prompt, optimizer, None, multi_apply_fn
-            )
-            logging.info(f"Initial rewards: {init_rewards}")
-            logging.info(f"Best rewards: {best_rewards}")
-            outpath = f"{outdir}/{index:0>5}"
-            os.makedirs(f"{outpath}/samples", exist_ok=True)
-            with open(f"{outpath}/metadata.jsonl", "w") as fp:
-                json.dump(metadata, fp)
-            best_image.save(f"{outpath}/samples/{args.seed:05}.png")
-            if i == 0:
-                total_best_rewards = {k: 0.0 for k in best_rewards.keys()}
-                total_init_rewards = {k: 0.0 for k in best_rewards.keys()}
-            for k in best_rewards.keys():
-                total_best_rewards[k] += best_rewards[k]
-                total_init_rewards[k] += init_rewards[k]
-        for k in total_best_rewards.keys():
-            total_best_rewards[k] /= len(parti_dataset)
-            total_init_rewards[k] /= len(parti_dataset)
+    except torch.cuda.OutOfMemoryError as e:
+        print(f"Global CUDA Out of Memory Error: {e}")
+        yield (None, f"{e}", None)
+    except RuntimeError as e:
+        if 'out of memory' in str(e):
+            print(f"Runtime Error: {e}")
+            yield (None, f"{e}", 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):
+    if gallery_state is not None:
+        return gr.update(value=gallery_state, visible=True)
     else:
-        raise ValueError(f"Unknown task {args.task}")
-    # log total rewards
-    logging.info(f"Mean initial rewards: {total_init_rewards}")
-    logging.info(f"Mean best rewards: {total_best_rewards}")
+        return gr.update(value=None, visible=False)
+
+def combined_function(gallery_state, loaded_model_setup, prompt, chosen_model, seed, n_iter, enable_hps, hps_w, enable_imagereward, imgrw_w, enable_pickscore, pcks_w, enable_clip, clip_w, learning_rate, progress=gr.Progress(track_tqdm=True)):
+    # Step 1: Start Over
+    gallery_state, output_image, status, iter_gallery_update = start_over(gallery_state)
+    model_status = ""  # No model status yet
+    yield gallery_state, output_image, status, iter_gallery_update, loaded_model_setup, model_status
+
+    # Step 2: Setup the model
+    model_status, new_loaded_model_setup = None, None
+    for model_status, new_loaded_model_setup in setup_model(
+        loaded_model_setup, prompt, chosen_model, seed, n_iter, enable_hps, hps_w, 
+        enable_imagereward, imgrw_w, enable_pickscore, pcks_w, enable_clip, clip_w, learning_rate):
+        yield gallery_state, output_image, status, iter_gallery_update, new_loaded_model_setup, model_status
+
+    # Step 3: Generate the image
+    output_image, status, gallery_state_update = None, None, None
+    for output_image, status, gallery_state_update in generate_image(new_loaded_model_setup, n_iter):
+        yield gallery_state_update, output_image, status, iter_gallery_update, new_loaded_model_setup, model_status
+
+    # Step 4: Show the gallery
+    iter_gallery_update = show_gallery_output(gallery_state_update)
+    yield gallery_state_update, output_image, status, iter_gallery_update, new_loaded_model_setup, model_status
+
+
+# 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."
+
+css="""
+#model-status-id{
+    height: 126px;
+}
+#model-status-id .progress-text{
+    font-size: 10px!important;
+}
+#model-status-id .progress-level-inner{
+    font-size: 8px!important;
+}
+"""
+
+with gr.Blocks(css=css, analytics_enabled=False) as demo:
+    loaded_model_setup = gr.State()
+    gallery_state = gr.State()
+    with gr.Column():
+        gr.Markdown(title)
+        gr.Markdown(description)
+        gr.HTML("""
+        <div style="display:flex;column-gap:4px;">
+            <a href='https://github.com/ExplainableML/ReNO'>
+                <img src='https://img.shields.io/badge/GitHub-Repo-blue'>
+            </a> 
+            <a href='https://arxiv.org/abs/2406.04312v1'>
+                <img src='https://img.shields.io/badge/Paper-Arxiv-red'>
+            </a>
+        </div>
+        """)
+
+        with gr.Row():
+            with gr.Column():
+                prompt = gr.Textbox(label="Prompt")
+                with gr.Row():
+                    chosen_model = gr.Dropdown(["sd-turbo", "sdxl-turbo", "pixart", "hyper-sd", "flux"], label="Model", value="sd-turbo")
+                    seed = gr.Number(label="seed", value=0)
+
+                model_status = gr.Textbox(label="model status", visible=True, elem_id="model-status-id")
+                
+                with gr.Row():
+                    n_iter = gr.Slider(minimum=10, maximum=100, step=10, value=10, label="Number of Iterations")
+                    learning_rate = gr.Slider(minimum=0.1, maximum=10.0, step=0.1, value=5.0, label="Learning Rate")
+
+                with gr.Accordion("Advanced Settings", open=True):
+                    with gr.Column():
+                        with gr.Row():
+                            enable_hps = gr.Checkbox(label="HPS ON", value=False, scale=1)
+                            hps_w = gr.Slider(label="HPS weight", step=0.1, minimum=0.0, maximum=10.0, value=5.0, interactive=False, scale=3)
+                        with gr.Row():
+                            enable_imagereward = gr.Checkbox(label="ImageReward ON", value=False, scale=1)
+                            imgrw_w = gr.Slider(label="ImageReward weight", step=0.1, minimum=0, maximum=5.0, value=1.0, interactive=False, scale=3)
+                        with gr.Row():
+                            enable_pickscore = gr.Checkbox(label="PickScore ON", value=False, scale=1)
+                            pcks_w = gr.Slider(label="PickScore weight", step=0.01, minimum=0, maximum=5.0, value=0.05, interactive=False, scale=3)
+                        with gr.Row():
+                            enable_clip = gr.Checkbox(label="CLIP ON", value=False, scale=1)
+                            clip_w = gr.Slider(label="CLIP weight", step=0.01, minimum=0, maximum=0.1, value=0.01, interactive=False, scale=3)
+
+                submit_btn = gr.Button("Submit")
+
+                gr.Examples(
+                    examples = [
+                        "A red dog and a green cat",
+                        "A pink elephant and a grey cow",
+                        "A toaster riding a bike",
+                        "Dwayne Johnson depicted as a philosopher king in an academic painting by Greg Rutkowski",
+                        "A curious, orange fox and a fluffy, white rabbit, playing together in a lush, green meadow filled with yellow dandelions",
+                        "An epic oil painting: a red portal infront of a cityscape, a solitary figure, and a colorful sky over snowy mountains"
+                    ],
+                    inputs = [prompt]     
+                )
+            
+            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, visible=False)
+
+    def allow_weighting(weight_type):
+        if weight_type is True:
+            return gr.update(interactive=True)
+        else:
+            return gr.update(interactive=False)
+    
+    enable_hps.change(
+        fn = allow_weighting,
+        inputs = [enable_hps],
+        outputs = [hps_w],
+        queue = False
+    )
+    enable_imagereward.change(
+        fn = allow_weighting,
+        inputs = [enable_imagereward],
+        outputs = [imgrw_w],
+        queue = False
+    )
+    enable_pickscore.change(
+        fn = allow_weighting,
+        inputs = [enable_pickscore],
+        outputs = [pcks_w],
+        queue = False
+    )
+    enable_clip.change(
+        fn = allow_weighting,
+        inputs = [enable_clip],
+        outputs = [clip_w],
+        queue = False
+    )
 
-def main():
-    args = parse_args()
-    args, trainer, device, dtype, shape, enable_grad, settings, pipe = setup(args, loaded_model_setup=None)
-    execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pipe)
+    submit_btn.click(
+        fn = combined_function,
+        inputs = [
+            gallery_state, loaded_model_setup, prompt, chosen_model, seed, n_iter, 
+            enable_hps, hps_w, enable_imagereward, imgrw_w, enable_pickscore, 
+            pcks_w, enable_clip, clip_w, learning_rate
+        ],
+        outputs = [
+            gallery_state, output_image, status, iter_gallery, loaded_model_setup, model_status  # Ensure `model_status` is included in the outputs
+        ]
+    )
+    
+    """
+    submit_btn.click(
+        fn = start_over,
+        inputs =[gallery_state], 
+        outputs = [gallery_state, output_image, status, iter_gallery]  
+    ).then(
+        fn = setup_model,
+        inputs = [loaded_model_setup, prompt, chosen_model, seed, n_iter, enable_hps, hps_w, enable_imagereward, imgrw_w, enable_pickscore, pcks_w, enable_clip, clip_w, learning_rate],
+        outputs = [model_status, loaded_model_setup]  # Load the new setup into the state
+    ).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
+    )
+    """
 
-if __name__ == "__main__":
-    main()
+# Launch the app
+demo.queue().launch(show_error=True, show_api=False)