Update src/pipeline.py

src/pipeline.py

@@ -56,10 +56,10 @@ def load_pipeline() -> Pipeline:
         cache_dir, subfolder="transformer", torch_dtype=torch.bfloat16, quantization_config=config
     )
     text_encoder_2 = T5EncoderModel.from_pretrained(
-        model_id, subfolder="text_encoder_2", torch_dtype=torch.bfloat16
+        cache_dir, subfolder="text_encoder_2", torch_dtype=torch.bfloat16, quantization_config=config
     )
     text_encoder = CLIPTextModel.from_pretrained(
-        model_id, subfolder="text_encoder",torch_dtype=torch.bfloat16
+        cache_dir, subfolder="text_encoder",torch_dtype=torch.bfloat16, quantization_config=config
     )
     # vae=AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=dtype)
     pipeline = DiffusionPipeline.from_pretrained(
@@ -73,8 +73,8 @@ def load_pipeline() -> Pipeline:
     torch.backends.cudnn.benchmark = True
     torch.backends.cuda.matmul.allow_tf32 = True
     torch.cuda.set_per_process_memory_fraction(0.95)
-    pipeline.text_encoder.to(memory_format=torch.channels_last)
-    pipeline.transformer.to(memory_format=torch.channels_last)
+    # pipeline.text_encoder.to(memory_format=torch.channels_last)
+    # pipeline.transformer.to(memory_format=torch.channels_last)
     # torch.jit.enable_onednn_fusion(True)