Update src/pipeline.py

src/pipeline.py

@@ -13,7 +13,7 @@ import diffusers
 import gc
 from diffusers import FluxTransformer2DModel, DiffusionPipeline
 from torchao.quantization import quantize_,int8_weight_only
-
+HOME = os.environ["HOME"]
 os.environ["TOKENIZERS_PARALLELISM"] = "True"
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 Pipeline = None
@@ -30,7 +30,7 @@ def load_pipeline() -> Pipeline:
     dtype, device = torch.bfloat16, "cuda"
     vae = AutoencoderTiny.from_pretrained("aifeifei798/taef1", torch_dtype=dtype)
     quantize_(vae, int8_weight_only())
-    model = FluxTransformer2DModel.from_pretrained("slobers/transgender", torch_dtype=dtype, use_safetensors=False)
+    model = FluxTransformer2DModel.from_pretrained(f"{HOME}/.cache/huggingface/hub/models--slobers--transgender/snapshots/cb99836efa0ed55856970269c42fafdaa0e44c5d", torch_dtype=dtype, use_safetensors=False)
     pipeline = DiffusionPipeline.from_pretrained(ckpt_id, vae=vae, transformer=model, torch_dtype=dtype,)
     torch.backends.cudnn.benchmark = True
     torch.backends.cuda.matmul.allow_tf32 = True