Update app.py

app.py

@@ -14,11 +14,14 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 #torch._inductor.config.epilogue_fusion = False
 #torch._inductor.config.coordinate_descent_check_all_directions = True
 
-pipe = AuraFlowPipeline.from_pretrained(
+pipe_v1 = AuraFlowPipeline.from_pretrained(
 	"fal/AuraFlow",
     torch_dtype=torch.float16
 ).to("cuda")
-
+pipe = AuraFlowPipeline.from_pretrained(
+	"fal/AuraFlow-v0.2",
+    torch_dtype=torch.float16
+).to("cuda")
 #pipe.transformer.to(memory_format=torch.channels_last)
 #pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead", fullgraph=True)
 #pipe.transformer.to(memory_format=torch.channels_last)
@@ -31,22 +34,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, 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, progress=gr.Progress(track_tqdm=True), model_version="0.2"):
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
         
     generator = torch.Generator().manual_seed(seed)
-    
-    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] 
+    if(model_version == "0.1"):
+        image = 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] 
+    else:
+        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 image, seed
 
@@ -74,7 +87,7 @@ with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""
         # AuraFlow 0.1
-        Demo of the [AuraFlow 0.1](https://huggingface.co/fal/AuraFlow) 6.8B parameters open source diffusion transformer model
+        Demo of the [AuraFlow 0.2](https://huggingface.co/fal/AuraFlow) 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)]
         """)
         
@@ -94,6 +107,10 @@ with gr.Blocks(css=css) as demo:
 
         with gr.Accordion("Advanced Settings", open=False):
             
+            model_version = gr.Dropdown(
+                ["0.1", "0.2"], label="Model version", value="0.1"
+            )
+            
             negative_prompt = gr.Text(
                 label="Negative prompt",
                 max_lines=1,
@@ -157,7 +174,7 @@ 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],
+        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, model_version],
         outputs = [result, seed]
     )