add flux.1 kontext vae

app.py

@@ -1,7 +1,7 @@
 import spaces
 import gradio as gr
 import torch
-from diffusers import AutoencoderKL
+from diffusers import AutoencoderKL, AutoencoderDC
 from diffusers.utils.remote_utils import remote_decode
 import torchvision.transforms.v2 as transforms
 from torchvision.io import read_image
@@ -63,6 +63,8 @@ class VAETester:
             "FLUX.1": AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae").to(self.device),
             "CogView4-6B": AutoencoderKL.from_pretrained("THUDM/CogView4-6B", subfolder="vae").to(self.device),
             "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),
         }
         # Define the desired order of models
         order = [
@@ -78,6 +80,8 @@ class VAETester:
             "FLUX.1",
             #"FLUX.1 (remote)",
             "CogView4-6B",
+			"dc-ae-f32c32-sana-1.0",
+            "FLUX.1-Kontext",
         ]
 
         # Construct the vae_models dictionary in the specified order
@@ -142,7 +146,7 @@ class VAETester:
             results[name] = self.process_image(img, model_config, tolerance)
         return results
 
-@spaces.GPU(duration=15)
+@spaces.GPU(duration=20)
 def test_all_vaes(image_path: str, tolerance: float, img_size: int):
     """Gradio interface function to test all VAEs"""
     tester = VAETester(img_size=img_size)
@@ -198,7 +202,7 @@ with gr.Blocks(title="VAE Performance Tester", css=".monospace-text {font-family
             with gr.Row():
                 diff_gallery = gr.Gallery(label="Difference Maps", columns=4, height=512)
                 recon_gallery = gr.Gallery(label="Reconstructed Images", columns=4, height=512)
-            scores_output = gr.Textbox(label="Sum of differences (lower is better) | Processing time (lower is faster)", lines=10, elem_classes="monospace-text")
+            scores_output = gr.Textbox(label="Sum of differences (lower is better) | Processing time (lower is faster)", lines=12, elem_classes="monospace-text")
 
     if examples:
         with gr.Row():

requirements.txt

@@ -2,8 +2,5 @@ torch
 torchvision
 accelerate
 safetensors
-huggingface_hub
-spaces>=0.34.2
-gradio>=5.25.2
 torchao  # latest diffusers request this
 git+https://github.com/huggingface/diffusers@main