Update src/pipeline.py

src/pipeline.py

@@ -37,11 +37,10 @@ def load_pipeline() -> Pipeline:
     ).to(memory_format=torch.channels_last)
     
     vae = AutoencoderTiny.from_pretrained("RobertML/FLUX.1-schnell-vae_fx", revision="00c83cdfdfe46992eb0ed45921eee34261fcb56e", torch_dtype=dtype)
-    vae.encoder.load_state_dict(torch.load("encoder.pth"), strict=False)
-    vae.decoder.load_state_dict(torch.load("decoder.pth"), strict=False)
     
     path = os.path.join(HF_HUB_CACHE, "models--RobertML--FLUX.1-schnell-int8wo/snapshots/307e0777d92df966a3c0f99f31a6ee8957a9857a")
     model = FluxTransformer2DModel.from_pretrained(path, torch_dtype=dtype, use_safetensors=False).to(memory_format=torch.channels_last)
+    
     pipeline = FluxPipeline.from_pretrained(
         ckpt_id,
         vae=vae,
@@ -50,6 +49,9 @@ def load_pipeline() -> Pipeline:
         text_encoder_2=text_encoder_2,
         torch_dtype=dtype,
         ).to(device)
+    basepath = os.path.join(HF_HUB_CACHE, "models--manbeast3b--Flux.1.schnell-vae-kl-p10/snapshots/facb90ac7d8e13df9a8c177f18b0d450a3e1ed41")
+    pipeline.vae.encoder.load_state_dict(torch.load(os.path.join(basepath, "encoder.pth")), strict=False)
+    pipeline.vae.decoder.load_state_dict(torch.load(os.path.join(basepath, "decoder.pth")), strict=False)
     pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune")
     quantize_(pipeline.vae, int8_weight_only())
     for _ in range(3):