Update src/pipeline.py

src/pipeline.py

@@ -75,17 +75,23 @@ def load_pipeline() -> Pipeline:
     vae.encoder = E(16)
     vae.decoder = D(16)
 
-    def lsd(p, mod, pfx):
-        sd = torch.load(p, map_location="cpu", weights_only=True)
-        f_sd = {k.strip(pfx): v for k, v in sd.items() if k.strip(pfx) in mod.state_dict() and v.size() == mod.state_dict()[k.strip(pfx)].size()}
-        print(f"num keys: {len(f_sd)} of {len(mod.state_dict())} from {len(sd.items())}")
-        mod.load_state_dict(f_sd, strict=False)
-        mod.to(dtype=torch.bfloat16)
-
-    lsd("encoder.pth", vae.encoder, "encoder.")
-    lsd("decoder.pth", vae.decoder, "decoder.")
-    vae.encoder.requires_grad_(False)
+    encoder_path = "encoder.pth"
+    decoder_path = "decoder.pth"
+
+    if encoder_path is not None:
+        encoder_state_dict = torch.load(encoder_path, map_location="cpu", weights_only=True)
+        filtered_state_dict = {k.strip('encoder.'): v for k, v in encoder_state_dict.items() if k.strip('encoder.') in vae.encoder.state_dict() and v.size() == vae.encoder.state_dict()[k.strip('encoder.')].size()}
+        vae.encoder.load_state_dict(filtered_state_dict, strict=False)
+        vae.encoder.to(dtype=dtype)
+            
+    if decoder_path is not None:
+        decoder_state_dict = torch.load(decoder_path, map_location="cpu", weights_only=True)
+        filtered_state_dict = {k.strip('decoder.'): v for k, v in decoder_state_dict.items() if k.strip('decoder.') in vae.decoder.state_dict() and v.size() == vae.decoder.state_dict()[k.strip('decoder.')].size()}
+        vae.decoder.load_state_dict(filtered_state_dict, strict=False)
+        vae.decoder.to(dtype=dtype)
+
     vae.decoder.requires_grad_(False)
+    vae.encoder.requires_grad_(False)
     
     # quantize_(vae, int8_weight_only())
     quantizer = ModelQuantization(vae)