Update src/pipeline.py

src/pipeline.py

@@ -8,6 +8,12 @@ from torch import Generator
 from loss import SchedulerWrapper
 from utils import register_normal_pipeline, register_faster_forward, register_parallel_pipeline, seed_everything 
 from onediffx import compile_pipe, save_pipe, load_pipe
+from diffusers import BitsAndBytesConfig
+
+nf4_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+)
 
 def callback_dynamic_cfg(pipe, step_index, timestep, callback_kwargs):
   if step_index == int(pipe.num_timesteps * 0.78):
@@ -22,22 +28,23 @@ def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
     if not pipeline:
         pipeline = StableDiffusionXLPipeline.from_pretrained(
             "stablediffusionapi/newdream-sdxl-20",
-            torch_dtype=torch.float16,
+            # torch_dtype=torch.float16,
+            quantization_config=nf4_config,
         ).to("cuda")
 
     # Register optimizations for performance
     # register_parallel_pipeline(pipeline)
-    register_normal_pipeline(pipeline)
-    register_faster_forward(pipeline.unet)
+    # register_normal_pipeline(pipeline)
+    # register_faster_forward(pipeline.unet)
     
     pipeline.scheduler = SchedulerWrapper(DDIMScheduler.from_config(pipeline.scheduler.config))
 
-    # pipeline = compile_pipe(pipeline)
+    pipeline = compile_pipe(pipeline)
     load_pipe(pipeline, dir="/home/sandbox/.cache/huggingface/hub/models--RobertML--cached-pipe-02/snapshots/58d70deae87034cce351b780b48841f9746d4ad7")
 
     for _ in range(1):
         deepcache_output = pipeline(prompt="telestereography, unstrengthen, preadministrator, copatroness, hyperpersonal, paramountness, paranoid, guaniferous", output_type="pil", num_inference_steps=20)
-    # pipeline.scheduler.prepare_loss()
+    pipeline.scheduler.prepare_loss()
     for _ in range(2):
         pipeline(prompt="telestereography, unstrengthen, preadministrator, copatroness, hyperpersonal, paramountness, paranoid, guaniferous", output_type="pil", num_inference_steps=20)
     return pipeline