Update src/pipeline.py

src/pipeline.py

@@ -19,21 +19,24 @@ torch._dynamo.config.suppress_errors = True
 Pipeline = None
 ids = "slobers/Flux.1.Schnella"
 Revision = "e34d670e44cecbbc90e4962e7aada2ac5ce8b55b"
+ckpt_id = "manbeast3b/flux.1-schnell-full1"
+ckpt_revision = "cb1b599b0d712b9aab2c4df3ad27b050a27ec146"
+
 
 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, transformer=transformer, local_files_only=True, torch_dtype=torch.bfloat16,)
-    
+    pipeline = FluxPipeline.from_pretrained(ckpt_id, revision=ckpt_revision, transformer=transformer, local_files_only=True, torch_dtype=torch.bfloat16,)
+    pipeline.to("cuda")
     pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)
-    basepath = os.path.join(HF_HUB_CACHE, "models--manbeast3b--Flux.1.schnell_eagle5_1_0.1_unst_7_2k/snapshots/b7a5ce1313327009093d3178220267d0cf669b76")
-    basepath = os.path.join(HF_HUB_CACHE, "models--manbeast3b--Flux.1.schnell_eagle5_1_0.1_unst_8/snapshots/3666a458a53e7dc83adfecb0bf955a0b4d575843")
-    basepath = os.path.join(HF_HUB_CACHE, "models--manbeast3b--Flux.1.schnell_eagle5_1_0.1_unst_13/snapshots/b3bdda899cd1961ec9b97bffde3ded31afa73ce3")
-    basepath = os.path.join(HF_HUB_CACHE, "models--manbeast3b--Flux.1.schnell_eagle5_1_0.1_unst_10/snapshots/20e4cf6ce3cc658237dfd6aae1d5f14bc6b3d1a4")
-    pipeline.vae.encoder.load_state_dict(torch.load(os.path.join(basepath, "encoder.pth")), strict=False)
-    pipeline.vae.decoder.load_state_dict(torch.load(os.path.join(basepath, "decoder.pth")), strict=False)
+    # basepath = os.path.join(HF_HUB_CACHE, "models--manbeast3b--Flux.1.schnell_eagle5_1_0.1_unst_7_2k/snapshots/b7a5ce1313327009093d3178220267d0cf669b76")
+    # basepath = os.path.join(HF_HUB_CACHE, "models--manbeast3b--Flux.1.schnell_eagle5_1_0.1_unst_8/snapshots/3666a458a53e7dc83adfecb0bf955a0b4d575843")
+    # basepath = os.path.join(HF_HUB_CACHE, "models--manbeast3b--Flux.1.schnell_eagle5_1_0.1_unst_13/snapshots/b3bdda899cd1961ec9b97bffde3ded31afa73ce3")
+    # basepath = os.path.join(HF_HUB_CACHE, "models--manbeast3b--Flux.1.schnell_eagle5_1_0.1_unst_10/snapshots/20e4cf6ce3cc658237dfd6aae1d5f14bc6b3d1a4")
+    # pipeline.vae.encoder.load_state_dict(torch.load(os.path.join(basepath, "encoder.pth")), strict=False)
+    # pipeline.vae.decoder.load_state_dict(torch.load(os.path.join(basepath, "decoder.pth")), strict=False)
     quantize_(pipeline.vae, int8_weight_only())
-    pipeline.to("cuda")
+    
     for _ in range(3):
         pipeline(prompt="insensible, timbale, pothery, electrovital, actinogram, taxis, intracerebellar, centrodesmus", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
     return pipeline