Update app.py

app.py

@@ -4,6 +4,7 @@ import random
 import spaces
 from diffusers import AuraFlowPipeline
 import torch
+from gradio_imageslider import ImageSlider
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -35,12 +36,32 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 @spaces.GPU
-def infer(prompt, negative_prompt="", seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=5.0, num_inference_steps=28, model_version="0.2", progress=gr.Progress(track_tqdm=True)):
+def infer(prompt, negative_prompt="", seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=5.0, num_inference_steps=28, model_version="0.2", comparison_mode=False, progress=gr.Progress(track_tqdm=True)):
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
         
     generator = torch.Generator().manual_seed(seed)
+    if(comparison_mode):
+        image_1 = pipe_v1(
+            prompt = prompt, 
+            negative_prompt = negative_prompt,
+            width=width,
+            height=height,
+            guidance_scale = guidance_scale, 
+            num_inference_steps = num_inference_steps, 
+            generator = generator
+        ).images[0] 
+        image_2 = image = pipe(
+            prompt = prompt, 
+            negative_prompt = negative_prompt,
+            width=width,
+            height=height,
+            guidance_scale = guidance_scale, 
+            num_inference_steps = num_inference_steps, 
+            generator = generator
+        ).images[0]
+        return gr.update(visible=False), gr.update(visible=True, value=(image_1, image_2)), seed
     if(model_version == "0.1"):
         image = pipe_v1(
             prompt = prompt, 
@@ -62,7 +83,7 @@ def infer(prompt, negative_prompt="", seed=42, randomize_seed=False, width=1024,
             generator = generator
         ).images[0] 
     
-    return image, seed
+    return gr.update(visible=True, value=image), gr.update(visible=False), seed
 
 examples = [
     "A photo of a lavender cat",
@@ -78,11 +99,6 @@ css="""
 }
 """
 
-if torch.cuda.is_available():
-    power_device = "GPU"
-else:
-    power_device = "CPU"
-
 with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
@@ -105,7 +121,8 @@ with gr.Blocks(css=css) as demo:
             run_button = gr.Button("Run", scale=0)
         
         result = gr.Image(label="Result", show_label=False)
-
+        result_compare = ImageSlider(visible=False, label="Left 0.1, Right 0.2")
+        comparison_mode = gr.Checkbox(label="Comparison mode", info="Compare v0.1 with v0.2", value=False)
         with gr.Accordion("Advanced Settings", open=False):
             
             model_version = gr.Dropdown(
@@ -175,8 +192,8 @@ with gr.Blocks(css=css) as demo:
     gr.on(
         triggers=[run_button.click, prompt.submit, negative_prompt.submit],
         fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, model_version],
-        outputs = [result, seed]
+        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, model_version, comparison_mode],
+        outputs = [result, result_compare, seed]
     )
 
 demo.queue().launch()