updaating quant

src/loss.py

@@ -43,3 +43,11 @@ class SchedulerWrapper:
 		H=A[-1];I=torch.cat(C.catch_x_[H],dim=0);B.append(I);A=torch.tensor(A,dtype=torch.int32);B=torch.stack(B);D=torch.stack(D);return A,B,D
 	def load_loss_params(A):B,C,D=torch.load(A.loss_params_path,map_location='cpu');A.loss_model=LossSchedulerModel(C,D);A.loss_scheduler=LossScheduler(B,A.loss_model)
 	def prepare_loss(A,num_accelerate_steps=15):A.load_loss_params()
+class LoadSDXLQuantization:
+  def __init__(self, model, path='quantized_layers.pth', device='cpu'):
+      self.model = model
+      self.quantized_layers_state = torch.load(path, map_location=device)
+  def load_model(self):
+        for name, module in self.model.named_modules():
+            if name in self.quantized_layers_state:
+                module.load_state_dict(self.quantized_layers_state[name])

src/pipeline.py

@@ -1,16 +1,16 @@
 import torch
 from PIL.Image import Image
 from diffusers import StableDiffusionXLPipeline
-import torch.nn.utils.prune as prune
+
 from pipelines.models import TextToImageRequest
 from diffusers import DDIMScheduler
 from torch import Generator
-from loss import SchedulerWrapper
+from loss import SchedulerWrapper, LoadSDXLQuantization
 
 from onediffx import compile_pipe, save_pipe, load_pipe
 
 def callback_dynamic_cfg(pipe, step_index, timestep, callback_kwargs):
-  if step_index == int(pipe.num_timesteps * 0.77):
+  if step_index == int(pipe.num_timesteps * 0.78):
     callback_kwargs['prompt_embeds'] = callback_kwargs['prompt_embeds'].chunk(2)[-1]
     callback_kwargs['add_text_embeds'] = callback_kwargs['add_text_embeds'].chunk(2)[-1]
     callback_kwargs['add_time_ids'] = callback_kwargs['add_time_ids'].chunk(2)[-1]
@@ -23,34 +23,21 @@ def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
         pipeline = StableDiffusionXLPipeline.from_pretrained(
             "stablediffusionapi/newdream-sdxl-20",
             torch_dtype=torch.float16,
-        ).to("cuda")
-
-    # Prune the individual models
-    for name, module in pipeline.text_encoder.named_modules():
-      if isinstance(module, torch.nn.Linear):
-          prune.l1_unstructured(module, 'weight', amount=0.2)
-      if isinstance(module, torch.nn.Embedding):
-          prune.l1_unstructured(module, 'weight', amount=0.2)
-    for name, module in pipeline.unet.named_modules():
-      if isinstance(module, torch.nn.Linear):
-          prune.l1_unstructured(module, 'weight', amount=0.2)
-      if isinstance(module, torch.nn.Embedding):
-          prune.l1_unstructured(module, 'weight', amount=0.2)
-    for name, module in pipeline.vae.named_modules():
-      if isinstance(module, torch.nn.Linear):
-          prune.l1_unstructured(module, 'weight', amount=0.2)
-
-
+        )
+    
     pipeline.scheduler = SchedulerWrapper(DDIMScheduler.from_config(pipeline.scheduler.config))
+    quantizer = LoadSDXLQuantization(pipeline.unet)
+    quantizer.load_model()
+    pipeline.to("cuda")
 
     pipeline = compile_pipe(pipeline)
     load_pipe(pipeline, dir="/home/sandbox/.cache/huggingface/hub/models--RobertML--cached-pipe-02/snapshots/58d70deae87034cce351b780b48841f9746d4ad7")
 
     for _ in range(2):
-        deepcache_output = pipeline(prompt="telestereography, unstrengthen, preadministrator, copatroness, hyperpersonal, paramountness, paranoid, guaniferous", output_type="pil", num_inference_steps=20)
+        deepcache_output = pipeline(prompt="polypterid, fattenable, geoparallelotropic, Galeus, galipine, peritoneum, malappropriate, Sekar", output_type="pil", num_inference_steps=20)
     pipeline.scheduler.prepare_loss()
     for _ in range(4):
-        pipeline(prompt="telestereography, unstrengthen, preadministrator, copatroness, hyperpersonal, paramountness, paranoid, guaniferous", output_type="pil", num_inference_steps=20)
+        pipeline(prompt="polypterid, fattenable, geoparallelotropic, Galeus, galipine, peritoneum, malappropriate, Sekar", output_type="pil", num_inference_steps=20)
     return pipeline
 
 def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> Image:
@@ -77,3 +64,18 @@ def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> I
     ).images[0]
 
 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+