enable TinyAE

app.py

@@ -56,7 +56,7 @@ class VAETester:
 			# "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": AutoModel.from_pretrained("fal/FLUX.2-Tiny-AutoEncoder", trust_remote_code=True, torch_dtype=torch.bfloat16).to(self.device),
+            "FLUX.2-TinyAutoEncoder": AutoModel.from_pretrained("fal/FLUX.2-Tiny-AutoEncoder", trust_remote_code=True, torch_dtype=torch.bfloat16).to(self.device),
         }
         # Define the desired order of models
         order = [
@@ -72,31 +72,28 @@ class VAETester:
 			# "dc-ae-f32c32-sana-1.0",
             "FLUX.1-Kontext",
             "FLUX.2",
-            # "FLUX.2-TinyAutoEncoder",
+            "FLUX.2-TinyAutoEncoder",
         ]
 
         # Construct the vae_models dictionary in the specified order
         return {name: {"vae": vaes[name], "dtype": torch.bfloat16 if name == "FLUX.2-TinyAutoEncoder" else torch.float32} for name in order}
 
-    def process_image(self, img: torch.Tensor, model_config: Dict, tolerance: float):
+    def process_image(self, img: torch.Tensor, model_config: Dict, tolerance: float, vae_name: str):
         """Process image through a single VAE model"""
         dtype = model_config["dtype"]
+        vae = model_config["vae"]
         img_transformed = self.input_transform(img).to(dtype).to(self.device).unsqueeze(0)
         original_base = self.base_transform(img).cpu()
 
-        # Start timer
+        # Time the encoding-decoding process
         start_time = time.time()
-
-        vae = model_config["vae"]
         with torch.no_grad():
-            if isinstance(vae, AutoModel):
+            if vae_name == "FLUX.2-TinyAutoEncoder":
                 encoded = vae.encode(img_transformed, return_dict=False)
                 decoded = vae.decode(encoded, return_dict=False)
             else:
                 encoded = vae.encode(img_transformed).latent_dist.sample()
                 decoded = vae.decode(encoded).sample
-
-        # End timer
         processing_time = time.time() - start_time
 
         decoded_transformed = self.output_transform(decoded.squeeze(0).to(torch.float32)).cpu()
@@ -111,8 +108,8 @@ class VAETester:
     def process_all_models(self, img: torch.Tensor, tolerance: float):
         """Process image through all configured VAEs"""
         results = {}
-        for name, model_config in self.vae_models.items():
-            results[name] = self.process_image(img, model_config, tolerance)
+        for vae_name, model_config in self.vae_models.items():
+            results[vae_name] = self.process_image(img, model_config, tolerance, vae_name)
         return results
 
 @spaces.GPU(duration=20)