Update src/pipeline.py

src/pipeline.py

@@ -16,7 +16,9 @@ from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
 import os
 os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
 torch._dynamo.config.suppress_errors = True
-
+torch.backends.cudnn.benchmark = True
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.cuda.set_per_process_memory_fraction(0.99)
 Pipeline = None
 
 ckpt_id = "black-forest-labs/FLUX.1-schnell"
@@ -35,26 +37,30 @@ def load_pipeline() -> Pipeline:
     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)
-    
     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)
+    transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=dtype, use_safetensors=False).to(memory_format=torch.channels_last)
     pipeline = FluxPipeline.from_pretrained(
         ckpt_id,
         revision=ckpt_revision,
-        transformer=model,
+        transformer=transformer,
         text_encoder_2=text_encoder_2,
         torch_dtype=dtype,
         ).to(device)
     pipeline.transformer = torch.compile(pipeline.transformer, mode="reduce-overhead")
     quantize_(pipeline.vae, int8_weight_only())
-    for _ in range(5):
+    
+    for _ in range(4):
         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
 
-
-@torch.no_grad()
+sample = True
+@torch.inference_mode()
 def infer(request: TextToImageRequest, pipeline: Pipeline, generator: Generator) -> Image:
-   image=pipeline(request.prompt,generator=generator, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256, height=request.height, width=request.width, output_type="pil").images[0]
-   return(image)
+    global sample
+    if sample:
+        empty_cache()
+        sample = None
+    image=pipeline(request.prompt,generator=generator, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256, height=request.height, width=request.width, output_type="pil").images[0]
+    return(image)