Update src/pipeline.py

src/pipeline.py

@@ -113,16 +113,16 @@ def load_pipeline() -> Pipeline:
     model_name = "manbeast3b/Flux.1.schnell-quant2"
     text_enc_revision = "44eb293715147878512da10bf3bc47cd14ec8c55"
     path = os.path.join(HF_HUB_CACHE, "models--manbeast3b--Flux.1.schnell-quant2/snapshots/44eb293715147878512da10bf3bc47cd14ec8c55/transformer")
-    # transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=True)
-
-    # model_name = "manbeast3b/Flux.1.Schnell-full-quant1"
-    # text_enc_revision = "e7ddf488a4ea8a3cba05db5b8d06e7e0feb826a2"
-    # path = os.path.join(HF_HUB_CACHE, "models--manbeast3b--Flux.1.Schnell-full-quant1/snapshots/e7ddf488a4ea8a3cba05db5b8d06e7e0feb826a2/transformer")
-    # transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=True)
+    transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=False)
 
     # model_name = "manbeast3b/flux.1-schnell-full1"
     # text_enc_revision = "cb1b599b0d712b9aab2c4df3ad27b050a27ec146"
     # path = os.path.join(HF_HUB_CACHE, "models--manbeast3b--flux.1-schnell-full1/snapshots/cb1b599b0d712b9aab2c4df3ad27b050a27ec146/transformer")
+    # transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=False)
+
+    # model_name = "manbeast3b/Flux.1.Schnell-full-quant1"
+    # text_enc_revision = "e7ddf488a4ea8a3cba05db5b8d06e7e0feb826a2"
+    # path = os.path.join(HF_HUB_CACHE, "models--manbeast3b--Flux.1.Schnell-full-quant1/snapshots/e7ddf488a4ea8a3cba05db5b8d06e7e0feb826a2/transformer")
     # transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=True)
 
     # model_name = "manbeast3b/flux.1-schnell-full2"
@@ -132,12 +132,15 @@ def load_pipeline() -> Pipeline:
     # mainid = "manbeast3b/flux.1-schnell-full1"
     # mainrevision = "cb1b599b0d712b9aab2c4df3ad27b050a27ec146"
 
-    model_name = "manbeast3b/flux.1-schnell-full2"
-    text_enc_revision = "434e0ebe2c7d94a02bc45ddafcd6b2b2d7ae0ce9"
-    path = os.path.join(HF_HUB_CACHE, "models--manbeast3b--flux.1-schnell-full2/snapshots/434e0ebe2c7d94a02bc45ddafcd6b2b2d7ae0ce9/transformer")
-    transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=True)
-    mainid = "manbeast3b/Flux.1.schnell-quant2"
-    mainrevision = "44eb293715147878512da10bf3bc47cd14ec8c55"
+    # model_name = "manbeast3b/flux.1-schnell-full2"
+    # text_enc_revision = "434e0ebe2c7d94a02bc45ddafcd6b2b2d7ae0ce9"
+    # path = os.path.join(HF_HUB_CACHE, "models--manbeast3b--flux.1-schnell-full2/snapshots/434e0ebe2c7d94a02bc45ddafcd6b2b2d7ae0ce9/transformer")
+    # transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=True)
+    # mainid = "manbeast3b/Flux.1.schnell-quant2"
+    # mainrevision = "44eb293715147878512da10bf3bc47cd14ec8c55"
+
+    
+
 
     pipeline = FluxPipeline.from_pretrained(mainid, revision=mainrevision, transformer=transformer, local_files_only=True, torch_dtype=torch.bfloat16,)
     pipeline.to("cuda")