Update src/pipeline.py

src/pipeline.py

@@ -572,6 +572,42 @@ class FluxTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOrig
 
         return Transformer2DModelOutput(sample=output)
 
+def load_single_file_checkpoint(
+    pretrained_model_link_or_path,
+    force_download=False,
+    proxies=None,
+    token=None,
+    cache_dir=None,
+    local_files_only=None,
+    revision=None,
+):
+    import pdb; pdb.set_trace()
+    if os.path.isfile(pretrained_model_link_or_path):
+        pretrained_model_link_or_path = pretrained_model_link_or_path
+
+    else:
+        repo_id, weights_name = _extract_repo_id_and_weights_name(pretrained_model_link_or_path)
+        pretrained_model_link_or_path = _get_model_file(
+            repo_id,
+            weights_name=weights_name,
+            force_download=force_download,
+            cache_dir=cache_dir,
+            proxies=proxies,
+            local_files_only=local_files_only,
+            token=token,
+            revision=revision,
+        )
+    import pdb; pdb.set_trace()
+
+    checkpoint = load_state_dict(pretrained_model_link_or_path)
+
+    # some checkpoints contain the model state dict under a "state_dict" key
+    while "state_dict" in checkpoint:
+        checkpoint = checkpoint["state_dict"]
+
+    return checkpoint
+
+
 Pipeline = None
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.enabled = True
@@ -594,7 +630,15 @@ def load_pipeline() -> Pipeline:
     from diffusers.loaders.single_file_utils import create_diffusers_t5_model_from_checkpoint
     from diffusers.loaders.single_file_model import FromOriginalModelMixin
 
+    
     dtype, device = torch.bfloat16, "cuda"
+
+    import pdb; pdb.set_trace()
+
+    t5_path = os.path.join(HF_HUB_CACHE, "models--manbeast3b--t5-v1_1-xxl-encoder-q8/snapshots/59c6c9cb99dcea42067f32caac3ea0836ef4c548/t5-v1_1-xxl-encoder-Q8_0.gguf")
+    # config_path = os.path.join(HF_HUB_CACHE, "models--black-forest--labs/FLUX.1-schnell/snapshots/741f7c3ce8b383c54771c7003378a50191e9efe9/text_encoder_2/config.json")
+    config_path = os.path.join(HF_HUB_CACHE, "models--black-forest-labs--FLUX.1-schnell/snapshots/741f7c3ce8b383c54771c7003378a50191e9efe9/")
+    ckpt_t5 = load_single_file_checkpoint(t5_path,local_files_only=True)
     
     text_encoder_2 = T5EncoderModel.from_pretrained(
         "silentdriver/aadb864af9", revision = "060dabc7fa271c26dfa3fd43c16e7c5bf3ac7892", torch_dtype=torch.bfloat16