Update src/pipeline.py

src/pipeline.py

@@ -8,21 +8,19 @@ from diffusers import FluxPipeline, AutoencoderTiny
 
 Pipeline = None
 
-pipeline = FluxPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16
-)
-pipeline.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16)
-pipeline.enable_sequential_cpu_offload()
+
+def load_pipeline() -> Pipeline:
+    pipeline = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
+    pipeline.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16)
+    pipeline.enable_sequential_cpu_offload()
 
     for _ in range(2):
-        empty_cache()
         pipeline(prompt="insensible, timbale, pothery, electrovital, actinogram, taxis, intracerebellar, centrodesmus", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
-    return pipeline
-
 
+    return pipeline
 
 def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
-    empty_cache()
     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)
+    return(image)