Update src/pipeline.py

src/pipeline.py

@@ -577,10 +577,8 @@ torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.enabled = True
 torch.backends.cudnn.benchmark = True
 
-# ckpt_id = "black-forest-labs/FLUX.1-schnell"
-# ckpt_revision = "741f7c3ce8b383c54771c7003378a50191e9efe9" 
-ckpt_id = "silentdriver/4b68f38c0b"
-ckpt_revision = "36a3cf4a9f733fc5f31257099b56b304fb2eceab" 
+ckpt_id = "black-forest-labs/FLUX.1-schnell"
+ckpt_revision = "741f7c3ce8b383c54771c7003378a50191e9efe9" 
 def empty_cache():
     gc.collect()
     torch.cuda.empty_cache()
@@ -591,36 +589,41 @@ def load_pipeline() -> Pipeline:
     empty_cache()
 
     dtype, device = torch.bfloat16, "cuda"
+
+    text_encoder = CLIPTextModel.from_pretrained(
+        ckpt_id, subfolder="text_encoder", torch_dtype=torch.bfloat16
+    )
     
     text_encoder_2 = T5EncoderModel.from_pretrained(
         "city96/t5-v1_1-xxl-encoder-bf16", revision = "1b9c856aadb864af93c1dcdc226c2774fa67bc86", torch_dtype=torch.bfloat16
     ).to(memory_format=torch.channels_last)
-    
-    
+
+    text_encoder = CLIPTextModel.from_pretrained(
+        os.path.join(HF_HUB_CACHE, "models--manbeast3b--FLUX.1-schnell-te1/snapshots/05ac3e466d6b42b7794859560d875b25f6df5daf"), subfolder="text_encoder", torch_dtype=torch.bfloat16
+    ).to(memory_format=torch.channels_last)
+
     vae = AutoencoderTiny.from_pretrained("RobertML/FLUX.1-schnell-vae_e3m2", revision="da0d2cd7815792fb40d084dbd8ed32b63f153d8d", torch_dtype=dtype)
     
-    path = os.path.join(HF_HUB_CACHE, "models--RobertML--FLUX.1-schnell-int8wo/snapshots/307e0777d92df966a3c0f99f31a6ee8957a9857a")
+    path = os.path.join(HF_HUB_CACHE, "models--manbeast3b--FLUX.1-schnell-transformer-f8/snapshots/2ac0d29a2f3a00175fd638e82e8acaa4ddcbfd09")
     generator = torch.Generator(device=device)
     model = FluxTransformer2DModel.from_pretrained(path, torch_dtype=dtype, use_safetensors=False, generator= generator).to(memory_format=torch.channels_last)
     torch.backends.cudnn.benchmark = True
     torch.backends.cudnn.deterministic = False
-    # model = torch.compile(model, mode="max-autotune-no-cudagraphs")
-    # model = torch.compile(model,backend="aot_eager")
-    vae = torch.compile(vae)
+    
     pipeline = DiffusionPipeline.from_pretrained(
         ckpt_id,
+        text_encoder=text_encoder,
         vae=vae,
         revision=ckpt_revision,
         transformer=model,
         text_encoder_2=text_encoder_2,
         torch_dtype=dtype,
         ).to(device)
+    pipeline.vae = torch.compile(pipeline.vae)
     pipeline.vae.requires_grad_(False)
     pipeline.transformer.requires_grad_(False)
     pipeline.text_encoder_2.requires_grad_(False)
     pipeline.text_encoder.requires_grad_(False)
-
-    # pipeline.enable_sequential_cpu_offload(exclude=["transformer"])
     
     for _ in range(3):
         pipeline(prompt="blah blah waah waah oneshot oneshot gang gang", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)