Update src/pipeline.py

src/pipeline.py

@@ -29,61 +29,32 @@ def empty_cache():
     torch.cuda.reset_peak_memory_stats()
 
 def load_pipeline() -> Pipeline:    
-    # empty_cache()
+    empty_cache()
 
-    # dtype, device = torch.bfloat16, "cuda"
+    dtype, device = torch.bfloat16, "cuda"
     
-    # 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)
-    
-    # vae = AutoencoderTiny.from_pretrained("RobertML/FLUX.1-schnell-vae_fx", revision="00c83cdfdfe46992eb0ed45921eee34261fcb56e", torch_dtype=dtype)
-    # path = os.path.join(HF_HUB_CACHE, "models--RobertML--FLUX.1-schnell-int8wo/snapshots/307e0777d92df966a3c0f99f31a6ee8957a9857a")
-    # model = FluxTransformer2DModel.from_pretrained(path, torch_dtype=dtype, use_safetensors=False).to(memory_format=torch.channels_last)
-    # pipeline = FluxPipeline.from_pretrained(
-    #     ckpt_id,
-    #     vae=vae,
-    #     revision=ckpt_revision,
-    #     transformer=model,
-    #     text_encoder_2=text_encoder_2,
-    #     torch_dtype=dtype,
-    #     ).to(device)
-    # pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune")
-    # quantize_(pipeline.vae, int8_weight_only())
-    # for _ in range(3):
-    #     pipeline(prompt="onomancy, aftergo, spirantic, Platyhelmia, modificator, drupaceous, jobbernowl, hereness", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
-    
-    # empty_cache()
-
-
-    text_encoder = CLIPTextModel.from_pretrained(CHECKPOINT, revision=REVISION, subfolder="text_encoder", local_files_only=True, torch_dtype=torch.bfloat16,)
-    quantize(text_encoder, weights=qint8)
-    text_encoder_2 = T5EncoderModel.from_pretrained(CHECKPOINT, revision=REVISION, subfolder="text_encoder_2", local_files_only=True, torch_dtype=torch.bfloat16,)
-
-    path1 = os.path.join(HF_HUB_CACHE, "models--madebyollin--taef1/snapshots/5463ee684fd9131a724bea777a2f50d89b0b6b24")
-
-    vae = AutoencoderTiny.from_pretrained(path1, 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)
 
+    vae = AutoencoderTiny.from_pretrained(tinypath, torch_dtype=dtype)    
     path = os.path.join(HF_HUB_CACHE, "models--RobertML--FLUX.1-schnell-int8wo/snapshots/307e0777d92df966a3c0f99f31a6ee8957a9857a")
-
-    transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16, use_safetensors=False,)
-    transformer = torch.compile(transformer, mode="max-autotune", fullgraph=True)
-    quantize_(vae, int8_weight_only())
-
-    pipeline = FluxPipeline.from_pretrained(CHECKPOINT, revision=REVISION, local_files_only=True, text_encoder=text_encoder, text_encoder_2=text_encoder_2, transformer=transformer, vae=vae, torch_dtype=torch.bfloat16,).to("cuda")
-
-    generator = Generator(pipeline.device).manual_seed(42)
-    
+    model = FluxTransformer2DModel.from_pretrained(path, torch_dtype=dtype, use_safetensors=False).to(memory_format=torch.channels_last)
+    pipeline = FluxPipeline.from_pretrained(
+        ckpt_id,
+        vae=vae,
+        revision=ckpt_revision,
+        transformer=model,
+        # text_encoder_2=text_encoder_2,
+        torch_dtype=dtype,
+        ).to(device)
+    pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune")
+    quantize_(pipeline.text_encoder, int8_weight_only())
+    quantize_(pipeline.vae, int8_weight_only())
     for _ in range(3):
-        pipeline(
-            "intracerebellar, centrodesmus", 
-            generator=generator,
-            num_inference_steps=4,
-            max_sequence_length=256,
-            guidance_scale=0.0,
-            height=1024,
-            width=1024,
-        )
+        pipeline(prompt="onomancy, aftergo, spirantic, Platyhelmia, modificator, drupaceous, jobbernowl, hereness", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
+    
+    empty_cache()
     return pipeline