Update src/pipeline.py

src/pipeline.py

@@ -26,11 +26,15 @@ torch.backends.cudnn.benchmark = True
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.cuda.set_per_process_memory_fraction(0.95)
 
-def callback_dynamic_cfg(pipe, step_index, timestep, callback_kwargs):
-  if step_index == int(pipe.num_timesteps * 0.5):
-    callback_kwargs['prompt_embeds'] = callback_kwargs['prompt_embeds'].chunk(2)[-1]
-    pipe._guidance_scale = 0.0
-  return callback_kwargs
+def dynamic_range_quantization(self, bit_width=8):
+    """Dynamic range quantization with adaptive scaling"""
+
+    quantized_model = torch.quantization.quantize_dynamic(
+        self.model, 
+        {nn.Linear, nn.Conv2d}, 
+        dtype=torch.qint8 if bit_width == 8 else torch.quint8  # Dynamic dtype based on bit_width
+    )
+    return quantized_model
     
 class BasicQuantization:
     def __init__(self, bits=1):
@@ -105,8 +109,11 @@ def load_pipeline() -> Pipeline:
 
     ############ Text Encoder ############
     text_encoder = CLIPTextModel.from_pretrained(
-        ckpt_id, subfolder="text_encoder", torch_dtype=torch.bfloat16
+        ckpt_id, subfolder="text_encoder", #torch_dtype=torch.bfloat16
     )
+    text_encoder = dynamic_range_quantization(text_encodertext_encoder)
+    text_encoder.to(torch.bfloat16)
+    
     ############ Text Encoder 2 ############
     text_encoder_2 = T5EncoderModel.from_pretrained(
         "city96/t5-v1_1-xxl-encoder-bf16", torch_dtype=torch.bfloat16
@@ -149,8 +156,6 @@ def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
         max_sequence_length=256, 
         height=request.height, 
         width=request.width, 
-        # callback_on_step_end=callback_dynamic_cfg,
-        # callback_on_step_end_tensor_inputs=['prompt_embeds'],
         output_type="pt").images[0]
     
     return torchvision.transforms.functional.to_pil_image(image.to(torch.float32).mul_(2).sub_(1))