Update src/pipeline.py

src/pipeline.py

@@ -21,6 +21,8 @@ ids = "slobers/Flux.1.Schnella"
 Revision = "e34d670e44cecbbc90e4962e7aada2ac5ce8b55b"
 
 def load_pipeline() -> Pipeline:
+    path = os.path.join(HF_HUB_CACHE, "models--slobers--Flux.1.Schnella/snapshots/e34d670e44cecbbc90e4962e7aada2ac5ce8b55b/transformer")
+    transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=False)
     pipeline = FluxPipeline.from_pretrained(ids, revision=Revision, local_files_only=True, torch_dtype=torch.bfloat16,)
     pipeline.to("cuda")
     pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)