Update app.py

app.py

@@ -20,8 +20,13 @@ pipe_v1 = AuraFlowPipeline.from_pretrained(
     torch_dtype=torch.float16
 ).to("cuda")
 
+pipe_v2 = AuraFlowPipeline.from_pretrained(
+    "fal/AuraFlow-v0.2",
+    torch_dtype=torch.float16
+).to("cuda")
+
 pipe = AuraFlowPipeline.from_pretrained(
-	"fal/AuraFlow-v0.2",
+	"fal/AuraFlow-v0.3",
     torch_dtype=torch.float16
 ).to("cuda")
 #pipe.transformer.to(memory_format=torch.channels_last)
@@ -50,14 +55,25 @@ def infer_example(prompt, negative_prompt="", seed=42, randomize_seed=False, wid
     return image, seed
     
 @spaces.GPU(duration=95)
-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)):
+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.3",
+          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(
+        image_1 = pipe_v2(
             prompt = prompt, 
             negative_prompt = negative_prompt,
             width=width,
@@ -86,7 +102,17 @@ def infer(prompt, negative_prompt="", seed=42, randomize_seed=False, width=1024,
             guidance_scale = guidance_scale, 
             num_inference_steps = num_inference_steps, 
             generator = generator
-        ).images[0] 
+        ).images[0]
+    elif(model_version == "0.2"):
+        image = pipe_v2(
+            prompt = prompt, 
+            negative_prompt = negative_prompt,
+            width=width,
+            height=height,
+            guidance_scale = guidance_scale, 
+            num_inference_steps = num_inference_steps, 
+            generator = generator
+        ).images[0]
     else:
         image = pipe(
             prompt = prompt, 
@@ -96,8 +122,8 @@ def infer(prompt, negative_prompt="", seed=42, randomize_seed=False, width=1024,
             guidance_scale = guidance_scale, 
             num_inference_steps = num_inference_steps, 
             generator = generator
-        ).images[0] 
-    
+        ).images[0]
+        
     return gr.update(visible=True, value=image), gr.update(visible=False), seed
 
 examples = [
@@ -118,8 +144,8 @@ with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""
-        # AuraFlow 0.2
-        Demo of the [AuraFlow 0.2](https://huggingface.co/fal/AuraFlow) 6.8B parameters open source diffusion transformer model
+        # AuraFlow 0.3
+        Demo of the [AuraFlow 0.3](https://huggingface.co/fal/AuraFlow-v0.3) 6.8B parameters open source diffusion transformer model
         [[blog](https://blog.fal.ai/auraflow/)] [[model](https://huggingface.co/fal/AuraFlow)] [[fal](https://fal.ai/models/fal-ai/aura-flow)]
         """)
         
@@ -136,12 +162,12 @@ 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)
+        result_compare = ImageSlider(visible=False, label="Left 0.2, Right 0.3")
+        comparison_mode = gr.Checkbox(label="Comparison mode", info="Compare v0.2 with v0.3", value=False)
         with gr.Accordion("Advanced Settings", open=False):
             
             model_version = gr.Dropdown(
-                ["0.1", "0.2"], label="Model version", value="0.2"
+                ["0.1", "0.2", "0.3"], label="Model version", value="0.3"
             )
             
             negative_prompt = gr.Text(