Update src/pipeline.py

src/pipeline.py

@@ -136,42 +136,12 @@ def load_pipeline() -> Pipeline:
     return pipeline
 
 
-
-_inference_count = 0
-
 @torch.inference_mode()
 def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
-    global _inference_count
-    # Clear on first inference
-    if _inference_count == 0:
-        empty_cache()
-    
-    # Increment counter and empty cache every 4 inferences
-    _inference_count += 1
-    if _inference_count >= 4:
-        empty_cache()
-        _inference_count = 0
-    
-    torch.cuda.reset_peak_memory_stats()
-    generator = Generator("cuda").manual_seed(request.seed)
-    image = pipeline(
-            prompt=request.prompt,
-            generator=generator,
-            guidance_scale=0.0,
-            num_inference_steps=4,
-            max_sequence_length=256,
-            height=request.height,
-            width=request.width,
-            output_type="pt"
-        ).images[0]
-    return torchvision.transforms.functional.to_pil_image(image.to(torch.float32).mul_(2).sub_(1))
+    generator = Generator(pipeline.device).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]
 
-# @torch.inference_mode()
-# def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
-#     generator = Generator(pipeline.device).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)
     
-#     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="pt").images[0]
-#     return torchvision.transforms.functional.to_pil_image(image.to(torch.float32).mul_(2).sub_(1))
+    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="pt").images[0]
+    return torchvision.transforms.functional.to_pil_image(image.to(torch.float32).mul_(2).sub_(1))