Update src/pipeline.py

src/pipeline.py

@@ -35,6 +35,10 @@ def empty_cache():
     torch.cuda.reset_peak_memory_stats()
     print(f"Flush took: {time.time() - start}")
 
+
+cache_dir = "/root/.cache/huggingface/hub/models--manbeast3b--flux-schnell-int8/snapshots/eb656b7968de3088ccac7cda876f5782e5a2f721/"
+
+
 def load_pipeline() -> Pipeline:    
     empty_cache()
     dtype, device = torch.bfloat16, "cuda"
@@ -49,7 +53,7 @@ def load_pipeline() -> Pipeline:
 
     model_id = "manbeast3b/flux-schnell-int8"
     transformer = FluxTransformer2DModel.from_pretrained(
-        model_id, subfolder="transformer", torch_dtype=torch.bfloat16, local_files_only=True # quantization_config=config,
+        model_id, subfolder="transformer", torch_dtype=torch.bfloat16, cache_dir=cache_dir, # quantization_config=config,
     )
     text_encoder_2 = T5EncoderModel.from_pretrained(
         model_id, subfolder="text_encoder_2", torch_dtype=torch.bfloat16