Update src/pipeline.py

src/pipeline.py

@@ -14,6 +14,7 @@ import time
 from diffusers import FluxTransformer2DModel, DiffusionPipeline
 from torchao.quantization import quantize_, int8_weight_only
 #from torchao.quantization import autoquant
+from model import E, D
 Pipeline = None
 
 ckpt_id = "black-forest-labs/FLUX.1-schnell"
@@ -28,13 +29,45 @@ def empty_cache():
 def load_pipeline() -> Pipeline:    
     empty_cache()
 
+    vae = AutoencoderTiny.from_pretrained("madebyollin/taef1")
+    vae.encoder = E(16)
+    vae.decoder = D(16)
+    
+    e_p = "encoder.pth"
+    d_p = "decoder.pth"
+    
+    if e_p:
+        e_sd = torch.load(e_p, map_location="cpu", weights_only=True)
+        f_sd = {k.strip("encoder."): v for k, v in e_sd.items() if k.strip("encoder.") in vae.encoder.state_dict() and v.size() == vae.encoder.state_dict()[k.strip("encoder.")].size()}
+        print(f" num of keys in filtered: {len(f_sd)} and in decoder: {len(vae.encoder.state_dict())}")
+        vae.encoder.load_state_dict(f_sd, strict=False)
+    
+    if d_p:
+        d_sd = torch.load(d_p, map_location="cpu", weights_only=True)
+        f_sd = {k.strip("decoder."): v for k, v in d_sd.items() if k.strip("decoder.") in vae.decoder.state_dict() and v.size() == vae.decoder.state_dict()[k.strip("decoder.")].size()}
+        print(f" num of keys in filtered: {len(f_sd)} and in decoder: {len(vae.decoder.state_dict())}")
+        vae.decoder.load_state_dict(f_sd, strict=False)
+    
+    vae.decoder.requires_grad_(False)
+    vae.encoder.requires_grad_(False)
+    vae.to(dtype=DTYPE)
+
     dtype, device = torch.bfloat16, "cuda"
 
     empty_cache()
     pipeline = DiffusionPipeline.from_pretrained(
+        vae = vae,
         ckpt_id, 
         torch_dtype=dtype,
         )
+    
+    pipe.text_encoder.to(memory_format=torch.channels_last)
+    pipe.text_encoder_2.to(memory_format=torch.channels_last)
+    pipe.transformer.to(memory_format=torch.channels_last)
+    pipeline.vae.to(memory_format=torch.channels_last)
+    pipeline.vae = torch.compile(pipe.vae, mode="max-autotune")
+    pipeline._exclude_from_cpu_offload = ["vae"]
+
     pipeline.enable_sequential_cpu_offload()
     for _ in range(2):
         empty_cache()