Update src/pipeline.py

src/pipeline.py

@@ -86,8 +86,18 @@ def load_pipeline() -> Pipeline:
     pipeline.vae.to(memory_format=torch.channels_last)
     pipeline.vae = torch.compile(pipeline.vae)
     
-    pipeline._exclude_from_cpu_offload = ["vae", "transformer"]
-    pipeline.enable_sequential_cpu_offload()
+    pipeline._exclude_from_cpu_offload = ["vae"]
+    # pipeline.enable_sequential_cpu_offload()
+    def custom_cpu_offload(model, device, offload_buffers=True):
+        state_dict = model.state_dict()
+        filtered_state_dict = {k: v for k, v in state_dict.items() if isinstance(v, torch.Tensor)}
+        for name, param in filtered_state_dict.items():
+            param.data = param.to(device)
+
+    custom_cpu_offload(pipeline.text_encoder, "cpu")
+    custom_cpu_offload(pipeline.text_encoder_2, "cpu")
+    custom_cpu_offload(pipeline.transformer, "cpu")
+
     for _ in range(2):
         pipeline(prompt="onomancy, aftergo, spirantic, Platyhelmia, modificator, drupaceous, jobbernowl, hereness", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)