code optimization

free up cuda chache + gc + load models on cpu only when needed and unload after used + saving up on zerogpu duration

app.py

@@ -3,8 +3,9 @@ from diffusers import StableDiffusion3Pipeline, StableDiffusionPipeline, StableD
 import gradio as gr
 import os
 import random
-import numpy as np
+import numpy
 from PIL import Image
+import gc # free up memory
 import spaces
 
 HF_TOKEN = os.getenv("HF_TOKEN")  # login with hf read token to access sd gated models
@@ -17,61 +18,65 @@ else:
     print("Using CPU")
 
 
-MAX_SEED = np.iinfo(np.int32).max
+MAX_SEED = numpy.iinfo(numpy.int32).max
 
-# Initialize the pipelines for each sd model
+# Global dictionary to store pipelines
+PIPELINES = {}
 
-# sd3 medium
-sd3_medium_pipe = StableDiffusion3Pipeline.from_pretrained(
-    "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16
-)
-sd3_medium_pipe.enable_model_cpu_offload()
-
-# sd 2.1
-sd2_1_pipe = StableDiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16
-)
-sd2_1_pipe.enable_model_cpu_offload()
-
-# sdxl
-sdxl_pipe = StableDiffusionXLPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
-)
-sdxl_pipe.enable_model_cpu_offload()
 
-# sdxl flash
-sdxl_flash_pipe = StableDiffusionXLPipeline.from_pretrained(
-    "sd-community/sdxl-flash", torch_dtype=torch.float16
-)
-sdxl_flash_pipe.enable_model_cpu_offload()
-# Ensure sampler uses "trailing" timesteps for sdxl flash.
-sdxl_flash_pipe.scheduler = DPMSolverSinglestepScheduler.from_config(
-    sdxl_flash_pipe.scheduler.config, timestep_spacing="trailing"
-)
+def load_pipeline(model_choice):
+    """Loads the specified pipeline and stores it in the PIPELINES dictionary."""
+    if model_choice not in PIPELINES:
+        if model_choice == "sd3 medium":
+            PIPELINES[model_choice] = StableDiffusion3Pipeline.from_pretrained(
+                "stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16
+            )
+        elif model_choice == "sd2.1":
+            PIPELINES[model_choice] = StableDiffusionPipeline.from_pretrained(
+                "stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16
+            )
+        elif model_choice == "sdxl":
+            PIPELINES[model_choice] = StableDiffusionXLPipeline.from_pretrained(
+                "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
+            )
+        elif model_choice == "sdxl flash":
+            PIPELINES[model_choice] = StableDiffusionXLPipeline.from_pretrained(
+                "sd-community/sdxl-flash", torch_dtype=torch.float16
+            )
+            # Store the original scheduler for resetting 
+            PIPELINES[model_choice].original_scheduler = PIPELINES[model_choice].scheduler 
+        elif model_choice == "stable cascade":
+            # Store both prior and decoder pipelines under 'stable cascade'
+            PIPELINES[model_choice] = {
+                'prior': StableCascadePriorPipeline.from_pretrained(
+                    "stabilityai/stable-cascade-prior", variant="bf16", torch_dtype=torch.bfloat16
+                ),
+                'decoder': StableCascadeDecoderPipeline.from_pretrained(
+                    "stabilityai/stable-cascade", variant="bf16", torch_dtype=torch.float16
+                )
+            }
+        elif model_choice == "sd1.5":
+            PIPELINES[model_choice] = StableDiffusionPipeline.from_pretrained(
+                "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
+            )
+        else:
+            raise ValueError(f"Invalid model choice: {model_choice}")
 
-# stable cascade
-stable_cascade_prior_pipe = StableCascadePriorPipeline.from_pretrained(
-    "stabilityai/stable-cascade-prior", variant="bf16", torch_dtype=torch.bfloat16
-)
-stable_cascade_prior_pipe.enable_model_cpu_offload()
-stable_cascade_decoder_pipe = StableCascadeDecoderPipeline.from_pretrained(
-    "stabilityai/stable-cascade", variant="bf16", torch_dtype=torch.float16
-)
-stable_cascade_decoder_pipe.enable_model_cpu_offload()
+    return PIPELINES[model_choice]
 
-# sd 1.5
-sd1_5_pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
-)
-sd1_5_pipe.enable_model_cpu_offload()
 
-# empty cache to free up gpu memory before inference
-torch.cuda.empty_cache()
+def unload_pipeline(model_choice):
+    """Unloads the specified pipeline from the PIPELINES dictionary and frees GPU memory."""
+    if model_choice in PIPELINES:
+        del PIPELINES[model_choice]
+        
+    torch.cuda.empty_cache()
+    gc.collect()
 
-# Helper function to generate images for a single model
 @spaces.GPU(duration=80)
-def generate_single_image(
+def run_inference(
     prompt,
+    pipe,
     negative_prompt,
     num_inference_steps,
     guidance_scale,
@@ -79,71 +84,114 @@ def generate_single_image(
     width,
     seed,
     num_images_per_prompt,
-    model_choice,
-    generator,
     prior_num_inference_steps=None,
     prior_guidance_scale=None,
     decoder_num_inference_steps=None,
     decoder_guidance_scale=None,
 ):
-    # Select the correct pipeline based on the model choice
-    if model_choice == "sd3 medium":
-        pipe = sd3_medium_pipe
-    elif model_choice == "sd2.1":
-        pipe = sd2_1_pipe
-    elif model_choice == "sdxl":
-        pipe = sdxl_pipe
-    elif model_choice == "sdxl flash":
-        pipe = sdxl_flash_pipe
-    elif model_choice == "stable cascade":
-        pipe = stable_cascade_prior_pipe
-    elif model_choice == "sd1.5":
-        pipe = sd1_5_pipe
-    else:
-        raise ValueError(f"Invalid model choice: {model_choice}")
+    """Runs inference with the specified pipeline and parameters."""
+
+    # Enable CPU offloading only if a GPU is available, for saving up RAM
+    if torch.cuda.is_available():
+        if isinstance(pipe, dict):  # Special handling for stable cascade
+            pipe['prior'].enable_model_cpu_offload()
+            pipe['decoder'].enable_model_cpu_offload()
+        else:
+            pipe.enable_model_cpu_offload()
 
-    # stable cascade has 2 different type of pipelines
-    if model_choice == "stable cascade":
-        prior_output = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_inference_steps=prior_num_inference_steps,
-            guidance_scale=prior_guidance_scale,
-            height=height,
-            width=width,
-            generator=generator,
-            num_images_per_prompt=num_images_per_prompt,
+    # Reset the sampler if the model is NOT SDXL Flash
+    if model_choice != "sdxl flash" and "sdxl flash" in PIPELINES:
+        PIPELINES["sdxl flash"].scheduler = PIPELINES["sdxl flash"].original_scheduler
+
+    # Apply SDXL Flash sampler ONLY if model_choice is 'sdxl flash'
+    if model_choice == "sdxl flash":
+        pipe.scheduler = DPMSolverSinglestepScheduler.from_config(
+            pipe.scheduler.config, timestep_spacing="trailing"
         )
 
-        output = stable_cascade_decoder_pipe(
-            image_embeddings=prior_output.image_embeddings.to(torch.float16),
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_inference_steps=decoder_num_inference_steps,
-            guidance_scale=decoder_guidance_scale,
-        ).images
+    if seed == 0:
+        seed = random.randint(1, MAX_SEED)
+
+    generator = torch.Generator().manual_seed(seed)
 
-        # the rest of the models have similar pipeline
-    else:
-        output = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            height=height,
-            width=width,
-            generator=generator,
-            num_images_per_prompt=num_images_per_prompt,
-        ).images
 
-    # empty cache to free up gpu memory
-    torch.cuda.empty_cache()
+    if isinstance(pipe, dict): # Stable Cascade
+        with torch.inference_mode():
+            prior_output = pipe['prior'](
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                num_inference_steps=prior_num_inference_steps,
+                guidance_scale=prior_guidance_scale,
+                height=height,
+                width=width,
+                generator=generator,
+                num_images_per_prompt=num_images_per_prompt,
+            )
+        with torch.inference_mode():
+            output = pipe['decoder'](
+                image_embeddings=prior_output.image_embeddings.to(torch.float16),
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                num_inference_steps=decoder_num_inference_steps,
+                guidance_scale=decoder_guidance_scale,
+            ).images
+    else:  # Other pipelines
+        with torch.inference_mode():
+            output = pipe(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                height=height,
+                width=width,
+                generator=generator,
+                num_images_per_prompt=num_images_per_prompt,
+            ).images
 
     return output
 
+# Helper function to generate images for a single model
+def generate_single_image(
+    prompt,
+    model_choice,
+    negative_prompt,
+    num_inference_steps,
+    guidance_scale,
+    height,
+    width,
+    seed,
+    num_images_per_prompt,
+    prior_num_inference_steps=None,
+    prior_guidance_scale=None,
+    decoder_num_inference_steps=None,
+    decoder_guidance_scale=None,
+):
+    # Load the pipeline
+    pipe = load_pipeline(model_choice)
+
+    # Run inference
+    output = run_inference(
+        prompt,
+        pipe,
+        negative_prompt,
+        num_inference_steps,
+        guidance_scale,
+        height,
+        width,
+        seed,
+        num_images_per_prompt,
+        prior_num_inference_steps,
+        prior_guidance_scale,
+        decoder_num_inference_steps,
+        decoder_guidance_scale,
+    )
+
+    # Unload the pipeline
+    unload_pipeline(model_choice)
+
+    return output
 
 # Define the image generation function for the Arena tab
-@spaces.GPU(duration=240)
 def generate_arena_images(
     prompt,
     negative_prompt,
@@ -188,15 +236,12 @@ def generate_arena_images(
     decoder_guidance_scale_d,
     progress=gr.Progress(track_tqdm=True),
 ):
-    if seed == 0:
-        seed = random.randint(1, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
 
     # Generate images for selected models
     if num_models_to_compare >= 2:
         images_a = generate_single_image(
             prompt,
+            model_choice_a,
             negative_prompt,
             num_inference_steps_a,
             guidance_scale_a,
@@ -204,8 +249,6 @@ def generate_arena_images(
             width_a,
             seed,
             num_images_per_prompt,
-            model_choice_a,
-            generator,
             prior_num_inference_steps_a,
             prior_guidance_scale_a,
             decoder_num_inference_steps_a,
@@ -213,6 +256,7 @@ def generate_arena_images(
         )
         images_b = generate_single_image(
             prompt,
+            model_choice_b,
             negative_prompt,
             num_inference_steps_b,
             guidance_scale_b,
@@ -220,8 +264,6 @@ def generate_arena_images(
             width_b,
             seed,
             num_images_per_prompt,
-            model_choice_b,
-            generator,
             prior_num_inference_steps_b,
             prior_guidance_scale_b,
             decoder_num_inference_steps_b,
@@ -233,6 +275,7 @@ def generate_arena_images(
     if num_models_to_compare >= 3:
         images_c = generate_single_image(
             prompt,
+            model_choice_c,
             negative_prompt,
             num_inference_steps_c,
             guidance_scale_c,
@@ -240,8 +283,6 @@ def generate_arena_images(
             width_c,
             seed,
             num_images_per_prompt,
-            model_choice_c,
-            generator,
             prior_num_inference_steps_c,
             prior_guidance_scale_c,
             decoder_num_inference_steps_c,
@@ -253,6 +294,7 @@ def generate_arena_images(
     if num_models_to_compare >= 4:
         images_d = generate_single_image(
             prompt,
+            model_choice_d,
             negative_prompt,
             num_inference_steps_d,
             guidance_scale_d,
@@ -260,8 +302,6 @@ def generate_arena_images(
             width_d,
             seed,
             num_images_per_prompt,
-            model_choice_d,
-            generator,
             prior_num_inference_steps_d,
             prior_guidance_scale_d,
             decoder_num_inference_steps_d,
@@ -274,9 +314,9 @@ def generate_arena_images(
 
 
 # Define the image generation function for the Individual tab
-@spaces.GPU(duration=90)
 def generate_individual_image(
     prompt,
+    model_choice,
     negative_prompt,
     num_inference_steps,
     guidance_scale,
@@ -284,20 +324,16 @@ def generate_individual_image(
     width,
     seed,
     num_images_per_prompt,
-    model_choice,
     prior_num_inference_steps,
     prior_guidance_scale,
     decoder_num_inference_steps,
     decoder_guidance_scale,
     progress=gr.Progress(track_tqdm=True),
 ):
-    if seed == 0:
-        seed = random.randint(1, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
 
     output = generate_single_image(
         prompt,
+        model_choice,
         negative_prompt,
         num_inference_steps,
         guidance_scale,
@@ -305,8 +341,6 @@ def generate_individual_image(
         width,
         seed,
         num_images_per_prompt,
-        model_choice,
-        generator,
         prior_num_inference_steps,
         prior_guidance_scale,
         decoder_num_inference_steps,
@@ -630,18 +664,18 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         width_a = gr.Slider(
                             label="Width (Model A)",
                             info="Width of the Image",
-                            minimum=256,
-                            maximum=1344,
-                            step=32,
+                            minimum=512,
+                            maximum=1280,
                             value=1024,
+                            step=32,
                         )
                         height_a = gr.Slider(
                             label="Height (Model A)",
                             info="Height of the Image",
-                            minimum=256,
-                            maximum=1344,
-                            step=32,
+                            minimum=512,
+                            maximum=1280,
                             value=1024,
+                            step=32,
                         )
                     with gr.Column():
                         num_inference_steps_b = gr.Slider(
@@ -650,7 +684,7 @@ with gr.Blocks(theme=theme, css=css) as demo:
                             minimum=1,
                             maximum=50,
                             value=25,
-                            step=1,
+                            step=32,
                             visible=True,
                         )
                         guidance_scale_b = gr.Slider(
@@ -701,18 +735,18 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         width_b = gr.Slider(
                             label="Width (Model B)",
                             info="Width of the Image",
-                            minimum=256,
-                            maximum=1344,
-                            step=32,
+                            minimum=512,
+                            maximum=1280,
                             value=1024,
+                            step=32,
                         )
                         height_b = gr.Slider(
                             label="Height (Model B)",
                             info="Height of the Image",
-                            minimum=256,
-                            maximum=1344,
-                            step=32,
+                            minimum=512,
+                            maximum=1280,
                             value=1024,
+                            step=32,
                         )
                     with gr.Column(visible=False) as model_c_options:
                         num_inference_steps_c = gr.Slider(
@@ -772,18 +806,18 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         width_c = gr.Slider(
                             label="Width (Model C)",
                             info="Width of the Image",
-                            minimum=256,
-                            maximum=1344,
-                            step=32,
+                            minimum=512,
+                            maximum=1280,
                             value=1024,
+                            step=32,
                         )
                         height_c = gr.Slider(
                             label="Height (Model C)",
                             info="Height of the Image",
-                            minimum=256,
-                            maximum=1344,
-                            step=32,
+                            minimum=512,
+                            maximum=1280,
                             value=1024,
+                            step=32,
                         )
                     with gr.Column(visible=False) as model_d_options:
                         num_inference_steps_d = gr.Slider(
@@ -843,18 +877,18 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         width_d = gr.Slider(
                             label="Width (Model D)",
                             info="Width of the Image",
-                            minimum=256,
-                            maximum=1344,
-                            step=32,
+                            minimum=512,
+                            maximum=1280,
                             value=1024,
+                            step=32,
                         )
                         height_d = gr.Slider(
                             label="Height (Model D)",
                             info="Height of the Image",
-                            minimum=256,
-                            maximum=1344,
-                            step=32,
+                            minimum=512,
+                            maximum=1280,
                             value=1024,
+                            step=32,
                         )
                 with gr.Row():
                     seed = gr.Slider(
@@ -883,6 +917,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_a: gr.update(visible=True),
                         decoder_num_inference_steps_a: gr.update(visible=True),
                         decoder_guidance_scale_a: gr.update(visible=True),
+                        width_a: gr.update(step=512, value=1024, maximum=1536),
+                        height_a: gr.update(step=512, value=1024, maximum=1536),
                     }
                 elif model_choice_a == "sdxl flash":
                     return {
@@ -892,6 +928,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_a: gr.update(visible=False),
                         decoder_num_inference_steps_a: gr.update(visible=False),
                         decoder_guidance_scale_a: gr.update(visible=False),
+                        width_a: gr.update(step=32, value=1024, maximum=1536),
+                        height_a: gr.update(step=32, value=1024, maximum=1536),
                     }
                 elif model_choice_a == "sd1.5":
                     return {
@@ -900,26 +938,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_a: gr.update(visible=True),
                         decoder_num_inference_steps_a: gr.update(visible=True),
                         decoder_guidance_scale_a: gr.update(visible=True),
-                    }
-                elif model_choice_a == "sdxl flash":
-                    return {
-                        num_inference_steps_a: gr.update(visible=True, maximum=15, value=8),
-                        guidance_scale_a: gr.update(visible=True, maximum=6.0, value=3.5),
-                        prior_num_inference_steps_a: gr.update(visible=False),
-                        prior_guidance_scale_a: gr.update(visible=False),
-                        decoder_num_inference_steps_a: gr.update(visible=False),
-                        decoder_guidance_scale_a: gr.update(visible=False),
-                    }
-                elif model_choice_a == "sd1.5":
-                    return {
-                        num_inference_steps_a: gr.update(visible=True, maximum=50, value=25),
-                        guidance_scale_a: gr.update(visible=True, maximum=10.0, value=7.5),
-                        prior_num_inference_steps_a: gr.update(visible=False),
-                        prior_guidance_scale_a: gr.update(visible=False),
-                        decoder_num_inference_steps_a: gr.update(visible=False),
-                        decoder_guidance_scale_a: gr.update(visible=False),
-                        width_a: gr.update(value=512, maximum=768),
-                        height_a: gr.update(value=512, maximum=768),
+                        width_a: gr.update(step=32, value=512, maximum=768),
+                        height_a: gr.update(step=32, value=512, maximum=768),
                     }
                 elif model_choice_a == "sd2.1":
                     return {
@@ -929,8 +949,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_a: gr.update(visible=False),
                         decoder_num_inference_steps_a: gr.update(visible=False),
                         decoder_guidance_scale_a: gr.update(visible=False),
-                        width_a: gr.update(value=768, maximum=1024),
-                        height_a: gr.update(value=768, maximum=1024),
+                        width_a: gr.update(step=32, value=768, maximum=1024),
+                        height_a: gr.update(step=32, value=768, maximum=1024),
                     }
                 else:
                     return {
@@ -940,8 +960,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_a: gr.update(visible=False),
                         decoder_num_inference_steps_a: gr.update(visible=False),
                         decoder_guidance_scale_a: gr.update(visible=False),
-                        width_a: gr.update(maximum=1344),
-                        height_a: gr.update(maximum=1344),
+                        width_a: gr.update(step=32, value=1024, maximum=1536),
+                        height_a: gr.update(step=32, value=1024, maximum=1536),
                     }
 
             def toggle_visibility_arena_b(model_choice_b):
@@ -953,6 +973,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_b: gr.update(visible=True),
                         decoder_num_inference_steps_b: gr.update(visible=True),
                         decoder_guidance_scale_b: gr.update(visible=True),
+                        width_b: gr.update(step=256, value=1024, maximum=1536),
+                        height_b: gr.update(step=256, value=1024, maximum=1536),
                     }
                 elif model_choice_b == "sdxl flash":
                     return {
@@ -962,6 +984,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_b: gr.update(visible=False),
                         decoder_num_inference_steps_b: gr.update(visible=False),
                         decoder_guidance_scale_b: gr.update(visible=False),
+                        width_a: gr.update(step=32, value=1024, maximum=1536),
+                        height_a: gr.update(step=32, value=1024, maximum=1536),
                     }
                 elif model_choice_b == "sd1.5":
                     return {
@@ -971,8 +995,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_b: gr.update(visible=False),
                         decoder_num_inference_steps_b: gr.update(visible=False),
                         decoder_guidance_scale_b: gr.update(visible=False),
-                        width_b: gr.update(value=512, maximum=768),
-                        height_b: gr.update(value=512, maximum=768),
+                        width_b: gr.update(step=32, value=512, maximum=768),
+                        height_b: gr.update(step=32, value=512, maximum=768),
                     }
                 elif model_choice_b == "sd2.1":
                     return {
@@ -982,8 +1006,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_b: gr.update(visible=False),
                         decoder_num_inference_steps_b: gr.update(visible=False),
                         decoder_guidance_scale_b: gr.update(visible=False),
-                        width_b: gr.update(value=768, maximum=1024),
-                        height_b: gr.update(value=768, maximum=1024),
+                        width_b: gr.update(step=32, value=768, maximum=1024),
+                        height_b: gr.update(step=32, value=768, maximum=1024),
                     }
                 else:
                     return {
@@ -993,8 +1017,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_b: gr.update(visible=False),
                         decoder_num_inference_steps_b: gr.update(visible=False),
                         decoder_guidance_scale_b: gr.update(visible=False),
-                        width_b: gr.update(maximum=1344),
-                        height_b: gr.update(maximum=1344),
+                        width_b: gr.update(step=32, value=1024, maximum=1536),
+                        height_b: gr.update(step=32, value=1024, maximum=1536),
                     }
 
             def toggle_visibility_arena_c(model_choice_c):
@@ -1006,8 +1030,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_c: gr.update(visible=True),
                         decoder_num_inference_steps_c: gr.update(visible=True),
                         decoder_guidance_scale_c: gr.update(visible=True),
-                        width_c: gr.update(value=1024, maximum=1344),
-                        height_c: gr.update(value=1024, maximum=1344),
+                        width_c: gr.update(step=256, value=1024, maximum=1536),
+                        height_c: gr.update(step=256, value=1024, maximum=1536),
                     }
                 elif model_choice_c == "sdxl flash":
                     return {
@@ -1017,8 +1041,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_c: gr.update(visible=False),
                         decoder_num_inference_steps_c: gr.update(visible=False),
                         decoder_guidance_scale_c: gr.update(visible=False),
-                        width_c: gr.update(value=1024, maximum=1344),
-                        height_c: gr.update(value=1024, maximum=1344),
+                        width_c: gr.update(step=32, value=1024, maximum=1536),
+                        height_c: gr.update(step=32, value=1024, maximum=1536),
                     }
                 elif model_choice_c == "sd1.5":
                     return {
@@ -1028,8 +1052,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_c: gr.update(visible=False),
                         decoder_num_inference_steps_c: gr.update(visible=False),
                         decoder_guidance_scale_c: gr.update(visible=False),
-                        width_c: gr.update(value=512, maximum=768),
-                        height_c: gr.update(value=512, maximum=768),
+                        width_c: gr.update(step=32, value=512, maximum=768),
+                        height_c: gr.update(step=32, value=512, maximum=768),
                     }
                 elif model_choice_c == "sd2.1":
                     return {
@@ -1039,8 +1063,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_c: gr.update(visible=False),
                         decoder_num_inference_steps_c: gr.update(visible=False),
                         decoder_guidance_scale_c: gr.update(visible=False),
-                        width_c: gr.update(value=768, maximum=1024),
-                        height_c: gr.update(value=768, maximum=1024),
+                        width_c: gr.update(step=32, value=768, maximum=1024),
+                        height_c: gr.update(step=32, value=768, maximum=1024),
                     }
                 else:
                     return {
@@ -1050,8 +1074,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_c: gr.update(visible=False),
                         decoder_num_inference_steps_c: gr.update(visible=False),
                         decoder_guidance_scale_c: gr.update(visible=False),
-                        width_c: gr.update(value=1024, maximum=1344),
-                        height_c: gr.update(value=1024, maximum=1344),
+                        width_c: gr.update(step=32, value=1024, maximum=1536),
+                        height_c: gr.update(step=32, value=1024, maximum=1536),
                     }
 
             def toggle_visibility_arena_d(model_choice_d):
@@ -1063,8 +1087,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_d: gr.update(visible=True),
                         decoder_num_inference_steps_d: gr.update(visible=True),
                         decoder_guidance_scale_d: gr.update(visible=True),
-                        width_d: gr.update(value=1024, maximum=1344),
-                        height_d: gr.update(value=1024, maximum=1344),
+                        width_d: gr.update(step=256, value=1024, maximum=1536),
+                        height_d: gr.update(step=256, value=1024, maximum=1536),
                     }
                 elif model_choice_d == "sdxl flash":
                     return {
@@ -1074,8 +1098,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_d: gr.update(visible=False),
                         decoder_num_inference_steps_d: gr.update(visible=False),
                         decoder_guidance_scale_d: gr.update(visible=False),
-                        width_d: gr.update(value=1024, maximum=1344),
-                        height_d: gr.update(value=1024, maximum=1344),
+                        width_d: gr.update(step=32, value=1024, maximum=1536),
+                        height_d: gr.update(step=32, value=1024, maximum=1536),
                     }
                 elif model_choice_d == "sd1.5":
                     return {
@@ -1085,8 +1109,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_d: gr.update(visible=False),
                         decoder_num_inference_steps_d: gr.update(visible=False),
                         decoder_guidance_scale_d: gr.update(visible=False),
-                        width_d: gr.update(value=512, maximum=768),
-                        height_d: gr.update(value=512, maximum=768),
+                        width_d: gr.update(step=32, value=512, maximum=768),
+                        height_d: gr.update(step=32, value=512, maximum=768),
                     }
                 elif model_choice_d == "sd2.1":
                     return {
@@ -1096,8 +1120,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_d: gr.update(visible=False),
                         decoder_num_inference_steps_d: gr.update(visible=False),
                         decoder_guidance_scale_d: gr.update(visible=False),
-                        width_d: gr.update(value=768, maximum=1024),
-                        height_d: gr.update(value=768, maximum=1024),
+                        width_d: gr.update(step=32, value=768, maximum=1024),
+                        height_d: gr.update(step=32, value=768, maximum=1024),
                     }
                 else:
                     return {
@@ -1107,8 +1131,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale_d: gr.update(visible=False),
                         decoder_num_inference_steps_d: gr.update(visible=False),
                         decoder_guidance_scale_d: gr.update(visible=False),
-                        width_d: gr.update(value=1024, maximum=1344),
-                        height_d: gr.update(value=1024, maximum=1344),
+                        width_d: gr.update(step=32, value=1024, maximum=1536),
+                        height_d: gr.update(step=32, value=1024, maximum=1536),
                     }
 
             model_choice_a.change(
@@ -1402,18 +1426,18 @@ with gr.Blocks(theme=theme, css=css) as demo:
                     width = gr.Slider(
                         label="Width",
                         info="Width of the Image",
-                        minimum=256,
-                        maximum=1344,
-                        step=32,
+                        minimum=512,
+                        maximum=1280,
                         value=1024,
+                        step=32,
                     )
                     height = gr.Slider(
                         label="Height",
                         info="Height of the Image",
-                        minimum=256,
-                        maximum=1344,
-                        step=32,
+                        minimum=512,
+                        maximum=1280,
                         value=1024,
+                        step=32,
                     )
                 with gr.Row():
                     seed = gr.Slider(
@@ -1442,8 +1466,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale: gr.update(visible=True),
                         decoder_num_inference_steps: gr.update(visible=True),
                         decoder_guidance_scale: gr.update(visible=True),
-                        width: gr.update(value=1024, maximum=1344),
-                        height: gr.update(value=1024, maximum=1344),
+                        width: gr.update(step=256, value=1024, maximum=1536),
+                        height: gr.update(step=256, value=1024, maximum=1536),
                     }
                 elif model_choice == "sdxl flash":
                     return {
@@ -1453,8 +1477,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale: gr.update(visible=False),
                         decoder_num_inference_steps: gr.update(visible=False),
                         decoder_guidance_scale: gr.update(visible=False),
-                        width: gr.update(value=1024, maximum=1344),
-                        height: gr.update(value=1024, maximum=1344),
+                        width: gr.update(step=32, value=1024, maximum=1536),
+                        height: gr.update(step=32, value=1024, maximum=1536),
                     }
                 elif model_choice == "sd1.5":
                     return {
@@ -1464,8 +1488,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale: gr.update(visible=False),
                         decoder_num_inference_steps: gr.update(visible=False),
                         decoder_guidance_scale: gr.update(visible=False),
-                        width: gr.update(value=512, maximum=768),
-                        height: gr.update(value=512, maximum=768),
+                        width: gr.update(step=32, value=512, maximum=768),
+                        height: gr.update(step=32, value=512, maximum=768),
                     }
                 elif model_choice == "sd2.1":
                     return {
@@ -1475,8 +1499,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale: gr.update(visible=False),
                         decoder_num_inference_steps: gr.update(visible=False),
                         decoder_guidance_scale: gr.update(visible=False),
-                        width: gr.update(value=768, maximum=1024),
-                        height: gr.update(value=768, maximum=1024),
+                        width: gr.update(step=32, value=768, maximum=1024),
+                        height: gr.update(step=32, value=768, maximum=1024),
                     }
                 else:
                     return {
@@ -1486,8 +1510,8 @@ with gr.Blocks(theme=theme, css=css) as demo:
                         prior_guidance_scale: gr.update(visible=False),
                         decoder_num_inference_steps: gr.update(visible=False),
                         decoder_guidance_scale: gr.update(visible=False),
-                        width: gr.update(value=1024, maximum=1344),
-                        height: gr.update(value=1024, maximum=1344),
+                        width: gr.update(step=32, value=1024, maximum=1536),
+                        height: gr.update(step=32, value=1024, maximum=1536),
                     }
 
             model_choice.change(
@@ -1509,6 +1533,7 @@ with gr.Blocks(theme=theme, css=css) as demo:
                 examples=examples_individual,
                 inputs=[
                     prompt,
+                    model_choice,
                     negative_prompt,
                     num_inference_steps,
                     guidance_scale,
@@ -1516,7 +1541,6 @@ with gr.Blocks(theme=theme, css=css) as demo:
                     width,
                     seed,
                     num_images_per_prompt,
-                    model_choice,
                     prior_num_inference_steps,
                     prior_guidance_scale,
                     decoder_num_inference_steps,
@@ -1534,6 +1558,7 @@ with gr.Blocks(theme=theme, css=css) as demo:
                 fn=generate_individual_image,
                 inputs=[
                     prompt,
+                    model_choice,
                     negative_prompt,
                     num_inference_steps,
                     guidance_scale,
@@ -1541,7 +1566,6 @@ with gr.Blocks(theme=theme, css=css) as demo:
                     width,
                     seed,
                     num_images_per_prompt,
-                    model_choice,
                     prior_num_inference_steps,
                     prior_guidance_scale,
                     decoder_num_inference_steps,