Upload handler.py

handler.py

@@ -5,6 +5,19 @@ import io
 from PIL import Image
 from diffusers import AutoPipelineForImage2Image
 
+# torch==2.5.1+cu124
+# torchvision==0.18.0+cu124
+# torchaudio==2.5.1+cu124
+# diffusers==0.17.0.dev0
+# Pillow==10.0.0
+# fastapi
+# pydantic
+# uvicorn
+
+#torchvision==0.18.0
+# torchaudio==2.5.1
+
+
 class EndpointHandler:
     def __init__(self, path=""):
         """Initialize the model from the given path."""
@@ -48,6 +61,11 @@ class EndpointHandler:
             output_image.save(buffered, format="PNG")
             img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
             
+            # Clear any cache and memory used by the model after inference
+            torch.cuda.empty_cache()  # Clears GPU memory
+            del output_image  # Delete the output image from memory
+            del init_image  # Delete the input image from memory
+            
             return {"generated_image": img_str}
         
         except Exception as e: