Update src/pipeline.py

src/pipeline.py

@@ -1,4 +1,4 @@
-from diffusers import AutoencoderKL, AutoencoderTiny
+from diffusers import AutoencoderTiny
 from diffusers.image_processor import VaeImageProcessor
 import torch
 import torch._dynamo
@@ -7,7 +7,9 @@ from PIL.Image import Image
 from pipelines.models import TextToImageRequest
 from torch import Generator
 from diffusers import FluxPipeline
-from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
+from torchao.quantization import quantize_, int8_weight_only
+import os
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:False,garbage_collection_threshold:0.02"
 
 Pipeline = None
 MODEL_ID = "black-forest-labs/FLUX.1-schnell"
@@ -20,11 +22,7 @@ def clear():
 
 def load_pipeline() -> Pipeline:    
     clear()
-    # vae = AutoencoderKL.from_pretrained(
-    #     MODEL_ID, subfolder="vae", torch_dtype=torch.bfloat16
-    # )
     vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=DTYPE)
-    # quantize_(vae, fpx_weight_only(3, 2))
     quantize_(vae, int8_weight_only())
     pipeline = FluxPipeline.from_pretrained(MODEL_ID,vae=vae,
                                         torch_dtype=DTYPE)
@@ -43,18 +41,6 @@ def load_pipeline() -> Pipeline:
         pipeline(prompt="unpervaded, unencumber, froggish, groundneedle, transnatural, fatherhood, outjump, cinerator", width=1024, height=1024, guidance_scale=0.1, num_inference_steps=4, max_sequence_length=256)
     return pipeline
 
-# sample = True
-# @torch.inference_mode()
-# def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
-#     global sample
-#     if sample:
-#         clear()
-#         sample = None
-#     # torch.cuda.reset_peak_memory_stats()
-#     generator = Generator("cuda").manual_seed(request.seed)
-#     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)
-
 sample = True
 @torch.inference_mode()
 def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
@@ -62,9 +48,6 @@ def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
     if sample:
         clear()
         sample = None
-    # torch.cuda.reset_peak_memory_stats()
     generator = Generator("cuda").manual_seed(request.seed)
-    image = None
-    with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
-        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]
+    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)