Update src/pipeline.py

src/pipeline.py

@@ -8,12 +8,6 @@ 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):
@@ -28,7 +22,7 @@ def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
     if not pipeline:
         pipeline = StableDiffusionXLPipeline.from_pretrained(
             "stablediffusionapi/newdream-sdxl-20",
-            torch_dtype=torch.float16,
+            torch_dtype=torch.bfloat16,
         ).to("cuda")
 
     # Register optimizations for performance
@@ -42,12 +36,11 @@ def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
     # 
     load_pipe(pipeline, dir="/home/sandbox/.cache/huggingface/hub/models--RobertML--cached-pipe-02/snapshots/58d70deae87034cce351b780b48841f9746d4ad7")
 
-    for _ in range(2):
+    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()
-    for _ in range(4):
+    for _ in range(2):
         pipeline(prompt="telestereography, unstrengthen, preadministrator, copatroness, hyperpersonal, paramountness, paranoid, guaniferous", output_type="pil", num_inference_steps=20)
-    # save_pipe(pipeline, dir="/home/sandbox/.cache/pipe/")
     return pipeline
 
 def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> Image: