- add new AEs: Flux-2 and fal/TAE

app.py

@@ -65,6 +65,8 @@ class VAETester:
             "playground-v2.5": AutoencoderKL.from_pretrained("playgroundai/playground-v2.5-1024px-aesthetic", subfolder="vae").to(self.device),
 			# "dc-ae-f32c32-sana-1.0": AutoencoderDC.from_pretrained("mit-han-lab/dc-ae-f32c32-sana-1.0-diffusers").to(self.device),
             "FLUX.1-Kontext": AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", subfolder="vae").to(self.device),
+            "FLUX.2": AutoencoderKL.from_pretrained("black-forest-labs/FLUX.2-dev", subfolder="vae").to(self.device),
+            "FLUX.2-TinyAutoEncoder": AutoencoderKL.from_pretrained("fal/FLUX.2-Tiny-AutoEncoder").to(self.device),
         }
         # Define the desired order of models
         order = [
@@ -82,6 +84,8 @@ class VAETester:
             "CogView4-6B",
 			# "dc-ae-f32c32-sana-1.0",
             "FLUX.1-Kontext",
+            "FLUX.2",
+            "FLUX.2-TinyAutoEncoder",
         ]
 
         # Construct the vae_models dictionary in the specified order
@@ -170,9 +174,20 @@ def test_all_vaes(image_path: str, tolerance: float, img_size: int):
         return [None], [None], error_msg
 
 examples = [f"examples/{img_filename}" for img_filename in sorted(os.listdir("examples/"))]
-
-with gr.Blocks(title="VAE Performance Tester", css=".monospace-text {font-family: 'Courier New', Courier, monospace;}") as demo:
-    gr.Markdown("# VAE Comparison Tool")
+custom_css = """
+    .center-header {
+        display: flex;
+        align-items: center;
+        justify-content: center;
+        margin: 0 0 10px 0;
+    }
+    .monospace-text {
+        font-family: 'Courier New', Courier, monospace;
+    }
+"""
+
+with gr.Blocks(title="VAE Performance Tester", css=custom_css) as demo:
+    gr.Markdown("<div class='center-header'><h1>VAE Comparison Tool</h1></div>")
     gr.Markdown("""
         Upload an image or select an example to compare how different VAEs reconstruct it.
         1. The image is padded to a square and resized to the selected size (512 or 1024 pixels).