WIP: Modularization ("free(): invalid pointer" my ass)

1 files changed, 13 insertions, 2 deletions

diff --git a/infer.py b/infer.py
index 36b5a2c..2b07b21 100644
--- a/infer.py
+++ b/infer.py
@@ -214,10 +214,21 @@ def load_embeddings(pipeline, embeddings_dir):
 def create_pipeline(model, dtype):
     print("Loading Stable Diffusion pipeline...")
 
-    pipeline = VlpnStableDiffusion.from_pretrained(model, torch_dtype=dtype)
+    tokenizer = MultiCLIPTokenizer.from_pretrained(model, subfolder='tokenizer', torch_dtype=dtype)
+    text_encoder = CLIPTextModel.from_pretrained(model, subfolder='text_encoder', torch_dtype=dtype)
+    vae = AutoencoderKL.from_pretrained(model, subfolder='vae', torch_dtype=dtype)
+    unet = UNet2DConditionModel.from_pretrained(model, subfolder='unet', torch_dtype=dtype)
+    scheduler = DDIMScheduler.from_pretrained(model, subfolder='scheduler', torch_dtype=dtype)
 
-    patch_managed_embeddings(pipeline.text_encoder)
+    patch_managed_embeddings(text_encoder)
 
+    pipeline = VlpnStableDiffusion(
+        text_encoder=text_encoder,
+        vae=vae,
+        unet=unet,
+        tokenizer=tokenizer,
+        scheduler=scheduler,
+    )
     pipeline.enable_xformers_memory_efficient_attention()
     pipeline.enable_vae_slicing()
     pipeline.to("cuda")