Should address the failing tests

nod-ai · Feb 19, 2025 · 2899a70 · 2899a70
1 parent 84a2a3a
commit 2899a70
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 9 deletions.
diff --git a/sharktank/sharktank/models/vae/model.py b/sharktank/sharktank/models/vae/model.py
@@ -74,6 +74,15 @@ def forward(
                 "latent_embeds": latent_embeds,
             },
         )
+        if not self.hp.use_post_quant_conv:
+            sample = rearrange(
+                sample,
+                "b (h w) (c ph pw) -> b c (h ph) (w pw)",
+                h=math.ceil(1024 / 16),
+                w=math.ceil(1024 / 16),
+                ph=2,
+                pw=2,
+            )
         sample = sample / self.hp.scaling_factor + self.hp.shift_factor
 
         if self.hp.use_post_quant_conv:

diff --git a/sharktank/sharktank/pipelines/flux/export_components.py b/sharktank/sharktank/pipelines/flux/export_components.py
@@ -267,15 +267,7 @@ def __init__(self, weight_file, height=1024, width=1024, precision="fp32"):
         self.width = width
 
     def forward(self, z):
-        d_in = rearrange(
-            z,
-            "b (h w) (c ph pw) -> b c (h ph) (w pw)",
-            h=math.ceil(self.height / 16),
-            w=math.ceil(self.width / 16),
-            ph=2,
-            pw=2,
-        )
-        return self.ae.forward(d_in)
+        return self.ae.forward(z)
 
 
 def get_ae_model_and_inputs(