Update app.py

app.py

@@ -7,10 +7,10 @@ from torchao.quantization import autoquant
 
 
 # # normal FluxPipeline
-# pipeline_normal = FluxPipeline.from_pretrained(
-#     "sayakpaul/FLUX.1-merged",
-#     torch_dtype=torch.bfloat16
-# ).to("cuda")
+pipeline_normal = FluxPipeline.from_pretrained(
+    "sayakpaul/FLUX.1-merged",
+    torch_dtype=torch.bfloat16
+).to("cuda")
 
 # optimized FluxPipeline
 pipeline_optimized = FluxPipeline.from_pretrained(
@@ -38,7 +38,6 @@ for name, layer in pipeline_optimized.transformer.named_children():
 #     pipeline_optimized.transformer,
 #     error_on_unseen=False
 # )
-pipeline_normal = pipeline_optimized
 
 @spaces.GPU(duration=120)
 def generate_images(prompt, guidance_scale, num_inference_steps):
@@ -49,15 +48,14 @@ def generate_images(prompt, guidance_scale, num_inference_steps):
         num_inference_steps=int(num_inference_steps)
     ).images[0]
     
-    # # generate image with optimized pipeline
-    # image_optimized = pipeline_optimized(
-    #     prompt=prompt,
-    #     guidance_scale=guidance_scale,
-    #     num_inference_steps=int(num_inference_steps)
-    # ).images[0]
+    # generate image with optimized pipeline
+    image_optimized = pipeline_optimized(
+        prompt=prompt,
+        guidance_scale=guidance_scale,
+        num_inference_steps=int(num_inference_steps)
+    ).images[0]
     
-    # return image_normal, image_optimized
-    return image_normal
+    return image_normal, image_optimized
 
 # set up Gradio interface
 demo = gr.Interface(
@@ -69,7 +67,7 @@ demo = gr.Interface(
     ],
     outputs=[
         gr.Image(type="pil", label="Normal FluxPipeline"),
-        # gr.Image(type="pil", label="Optimized FluxPipeline")
+        gr.Image(type="pil", label="Optimized FluxPipeline")
     ],
     title="FluxPipeline Comparison",
     description="Compare images generated by the normal FluxPipeline and the optimized one using torchao and torch.compile()."