GPU management optimizations

app.py

@@ -1,372 +1,419 @@
+import json
+import logging
+import os
+
+import blobfile as bf
 import torch
 import gc
-import gradio as gr
-from main import setup, execute_task
+from datasets import load_dataset
+from pytorch_lightning import seed_everything
+from tqdm import tqdm
+
 from arguments import parse_args
-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.")
+from models import get_model, get_multi_apply_fn
+from rewards import get_reward_losses
+from training import LatentNoiseTrainer, get_optimizer
 
-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 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 ...")
-    
-    if prompt is None or prompt == "":
-        raise gr.Error("You forgot to provide a prompt !")
+import torch
+import gc
 
-    print(f"LOADED_MODEL SETUP: {loaded_model_setup}")
+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")
     
-    """Clear CUDA memory before starting the training."""
-    torch.cuda.empty_cache()  # Free up cached memory
+    # Force the garbage collector to free unreferenced objects
     gc.collect()
+
+    # 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
     
-    # 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
+    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}")
+    
+    # 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 ''}"
+    )
     
-    if enable_imagereward is True:
-        args.disable_imagereward = False
-        args.imagereward_weighting = imgrw_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_pickscore is True:
-        args.disable_pickscore = False
-        args.pickscore_weighting = pcks_w
+    logging.info(args)
     
-    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]
+    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
         
-        # 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'}
+        # 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 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
-            
-            return 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, multi_apply_fn, settings = setup(args, loaded_model_setup)
-        new_loaded_setup = [args, trainer, device, dtype, shape, enable_grad, multi_apply_fn, settings]
-        return f"{model} model loaded successfully!", new_loaded_setup
+        # 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,
+            )
+        
+        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}")
     
-    except Exception as e:
-        print(f"Failed to load {model} model: {e}.")
-        return f"Failed to load {model} model: {e}. You can try again, as it usually finally loads on the second try :)", None
-       
+    # Get reward losses
+    reward_losses = get_reward_losses(args, dtype, device, args.cache_dir)
 
-def generate_image(setup_args, num_iterations):
-    torch.cuda.empty_cache()  # Free up cached memory
+    # 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()
 
-    gr.Info(f"Executing iterations task ...")
+    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,
+    )
 
-    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]
-    multi_apply_fn = setup_args[6]
+    # 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,
+        )
+    
+    enable_grad = not args.no_optim
 
-    settings = setup_args[7]
-    print(f"SETTINGS: {settings}")
+    # Final memory cleanup
+    torch.cuda.empty_cache()  # Free up cached memory
+    gc.collect()
 
-    save_dir = f"{args.save_dir}/{args.task}/{settings}/{args.prompt[:150]}"
-    clean_dir(save_dir)
     
-    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
-        
-        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, multi_apply_fn, settings, 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()
+    return args, trainer, device, dtype, shape, enable_grad, settings, pipe
 
-    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)
+
+
+
+def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pipe, progress_callback=None):
+    
+    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:
-            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:
-        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."
-
-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]     
-                )
+            print(f"PIPE:{pipe}")
+            
             
-            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)
+            if args.cpu_offloading:
+                pipe.enable_sequential_cpu_offload()
+
+            #if pipe.device != torch.device('cuda'):
+            #    pipe.to(device, dtype)
+        
+        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:
-            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
-    )
+            multi_apply_fn = None    
     
+        torch.cuda.empty_cache()  # Free up cached memory
+        gc.collect()
     
-    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
-    )
 
-# Launch the app
-demo.queue().launch(show_error=True, show_api=False)
+        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)
+    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}")
+
+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)
+
+if __name__ == "__main__":
+    main()