manbeast3b
/

fire_stang3

Model card Files Files and versions Community

manbeast3b commited on 14 days ago

Commit

49e6e93

·

verified ·

1 Parent(s): 5396446

Update src/pipeline.py

Files changed (1) hide show

src/pipeline.py +7 -7

src/pipeline.py CHANGED Viewed

@@ -45,16 +45,16 @@ def _load_text_encoder_model():
 def _load_vae_model():
     """Load the variational autoencoder (VAE) model with specified configuration."""
-    # return AutoencoderTiny.from_pretrained(
-    #     "RobertML/FLUX.1-schnell-vae_e3m2",
-    #     revision="da0d2cd7815792fb40d084dbd8ed32b63f153d8d",
-    #     torch_dtype=Config.DTYPE
-    # )
     return AutoencoderTiny.from_pretrained(
-        "manbeast3b/FLUX.1-schnell-taef1-float8",
-        revision="7c538d53ec698509788ed88b1305c6bb019bdb4d",
         torch_dtype=Config.DTYPE
     )
 def _load_transformer_model():
     """Load the transformer model from a specific cached path."""

 def _load_vae_model():
     """Load the variational autoencoder (VAE) model with specified configuration."""
     return AutoencoderTiny.from_pretrained(
+        "RobertML/FLUX.1-schnell-vae_e3m2",
+        revision="da0d2cd7815792fb40d084dbd8ed32b63f153d8d",
         torch_dtype=Config.DTYPE
     )
+    # return AutoencoderTiny.from_pretrained(
+    #     "manbeast3b/FLUX.1-schnell-taef1-float8",
+    #     revision="7c538d53ec698509788ed88b1305c6bb019bdb4d",
+    #     torch_dtype=Config.DTYPE
+    # )
 def _load_transformer_model():
     """Load the transformer model from a specific cached path."""