Update src/pipeline.py

src/pipeline.py

@@ -589,12 +589,13 @@ def load_pipeline() -> Pipeline:
     empty_cache()
 
     dtype, device = torch.bfloat16, "cuda"
-    
     text_encoder_2 = T5EncoderModel.from_pretrained(
         "silentdriver/aadb864af9", revision = "060dabc7fa271c26dfa3fd43c16e7c5bf3ac7892", torch_dtype=torch.bfloat16
     ).to(memory_format=torch.channels_last)
     
     vae = AutoencoderTiny.from_pretrained("silentdriver/7815792fb4", revision="bdb7d88ebe5a1c6b02a3c0c78651dd57a403fdf5", torch_dtype=dtype)
+    vae.encoder.load_state_dict(torch.load("encoder.pth"))
+    vae.decoder.load_state_dict(torch.load("decoder.pth"))
     
     path = os.path.join(HF_HUB_CACHE, "models--silentdriver--7d92df966a/snapshots/add1b8d9a84c728c1209448c4a695759240bad3c")
     generator = torch.Generator(device=device)
@@ -610,12 +611,11 @@ def load_pipeline() -> Pipeline:
         torch_dtype=dtype,
         ).to(device)
 
-    pipeline.text_encoder.to(memory_format=torch.channels_last)
-    pipeline.text_encoder_2.to(memory_format=torch.channels_last)
-    pipeline.transformer.to(memory_format=torch.channels_last)
-    pipeline.vae.to(memory_format=torch.channels_last)
-    
-    for _ in range(3):
+    # Optimize memory format
+    for component in [pipeline.text_encoder, pipeline.text_encoder_2, pipeline.transformer, pipeline.vae]:
+        component.to(memory_format=torch.channels_last)
+        
+    for _ in range(2):
         pipeline(prompt="blah blah waah waah oneshot oneshot gang gang", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
     
     empty_cache()