Update src/pipeline.py

src/pipeline.py

@@ -11,34 +11,47 @@ from pipelines.models import TextToImageRequest
 from torch import Generator
 from diffusers import FluxTransformer2DModel, DiffusionPipeline
 from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
+from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
+os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
+os.environ["TOKENIZERS_PARALLELISM"] = "True"
+torch._dynamo.config.suppress_errors = True
 
 os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
 os.environ["TOKENIZERS_PARALLELISM"] = "True"
 torch._dynamo.config.suppress_errors = True
-Pipeline = None
-ids = "slobers/Flux.1.Schnella"
-Revision = "e34d670e44cecbbc90e4962e7aada2ac5ce8b55b"
+Pipeline: TypeAlias = FluxPipeline
+
+CHECKPOINT = "black-forest-labs/FLUX.1-schnell"
+REVISION = "741f7c3ce8b383c54771c7003378a50191e9efe9"
 
 
 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.to("cuda")
+    text_encoder = CLIPTextModel.from_pretrained(CHECKPOINT, revision=REVISION, subfolder="text_encoder", local_files_only=True, torch_dtype=torch.bfloat16,)
+    text_encoder_2 = T5EncoderModel.from_pretrained(CHECKPOINT, revision=REVISION, subfolder="text_encoder_2", local_files_only=True, torch_dtype=torch.bfloat16,)
+    vae = AutoencoderKL.from_pretrained(CHECKPOINT, revision=REVISION, subfolder="vae", local_files_only=True, torch_dtype=torch.bfloat16,)
+    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,)
+    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")
+    quantize_(pipeline.text_encoder, float8_weight_only())
     quantize_(pipeline.vae, int8_weight_only())
-    pipeline.vae = torch.compile(pipeline.vae, options={"epilogue_fusion": True, "max-autotune": True}, fullgraph=True, dynamic=True)
-    for _ in range(3):
-        pipeline(prompt="insensible, timbale, pothery, electrovital, actinogram, taxis, intracerebellar, centrodesmus", width=1024, height=1024, guidance_scale=5.0, num_inference_steps=4, max_sequence_length=256)
+    #pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune-no-cudagraphs",)
+    pipeline = apply_cache_on_pipe(pipeline, residual_diff_threshold=0.31)
+    
+    pipeline("")
+
     return pipeline
 
-@torch.no_grad()
 def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+
     generator = Generator(pipeline.device).manual_seed(request.seed)
 
     return pipeline(
         request.prompt,
         generator=generator,
-        guidance_scale=1.5,
+        guidance_scale=0.0,
         num_inference_steps=4,
         max_sequence_length=256,
         height=request.height,