Update src/pipeline.py

src/pipeline.py

@@ -45,16 +45,8 @@ def load_pipeline():
 
 
 @torch.inference_mode()
-def infer(request: TextToImageRequest, pipeline):
-    generator = torch.Generator("cuda").manual_seed(request.seed)
-    image = pipeline(
-        request.prompt,
-        generator=generator,
-        guidance_scale=0.0,  # Match original repos except -25, test higher values
-        num_inference_steps=4,
-        max_sequence_length=256,
-        height=request.height,
-        width=request.width,
-        output_type="pil"
-    ).images[0]
-    return image
+def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image:
+    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)