test

app.py

@@ -1,5 +1,5 @@
 import torch
-from diffusers import StableDiffusion3Pipeline, StableDiffusionPipeline, StableDiffusionXLPipeline, DPMSolverSinglestepScheduler
+from diffusers import StableDiffusion3Pipeline, StableDiffusionPipeline, DiffusionPipeline, StableDiffusionXLPipeline, DPMSolverSinglestepScheduler
 import gradio as gr
 import os
 import random
@@ -29,17 +29,57 @@ sd2_1_pipe = StableDiffusionPipeline.from_pretrained(
 )
 sd2_1_pipe.to(device)
 
-sdxl_pipe = StableDiffusionXLPipeline.from_pretrained(
+sdxl_pipe = DiffusionPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
 )
 sdxl_pipe.to(device)
 
-sdxl_flash_pipe = StableDiffusionXLPipeline.from_pretrained("sd-community/sdxl-flash", torch_dtype=torch.float16)
+sdxl_flash_pipe = StableDiffusionXLPipeline.from_pretrained(
+    "sd-community/sdxl-flash", torch_dtype=torch.float16
+)
 sdxl_flash_pipe.to(device)
-
 # Ensure sampler uses "trailing" timesteps for sdxl flash.
 sdxl_flash_pipe.scheduler = DPMSolverSinglestepScheduler.from_config(sdxl_flash_pipe.scheduler.config, timestep_spacing="trailing")
 
+# Helper function to generate images for a single model
+@spaces.GPU(duration=80)
+def generate_single_image(
+    prompt,
+    negative_prompt,
+    num_inference_steps,
+    height,
+    width,
+    guidance_scale,
+    seed,
+    num_images_per_prompt,
+    model_choice,
+    generator,
+):
+    # 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
+    else:
+        raise ValueError(f"Invalid model choice: {model_choice}")
+
+    output = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        num_inference_steps=num_inference_steps,
+        height=height,
+        width=width,
+        guidance_scale=guidance_scale,
+        generator=generator,
+        num_images_per_prompt=num_images_per_prompt,
+    ).images
+
+    return output
+
 # Define the image generation function for the Arena tab
 @spaces.GPU(duration=80)
 def generate_arena_images(
@@ -88,45 +128,6 @@ def generate_arena_images(
 
     return images_1, images_2
 
-
-# Helper function to generate images for a single model
-def generate_single_image(
-    prompt,
-    negative_prompt,
-    num_inference_steps,
-    height,
-    width,
-    guidance_scale,
-    seed,
-    num_images_per_prompt,
-    model_choice,
-    generator,
-):
-    # 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
-    else:
-        raise ValueError(f"Invalid model choice: {model_choice}")
-
-    output = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        num_inference_steps=num_inference_steps,
-        height=height,
-        width=width,
-        guidance_scale=guidance_scale,
-        generator=generator,
-        num_images_per_prompt=num_images_per_prompt,
-    ).images
-
-    return output
-
 # Define the image generation function for the Individual tab
 @spaces.GPU(duration=80)
 def generate_individual_image(