import runpod
from runpod.serverless.utils.rp_cleanup import clean
import requests
import os
def download_image(url, save_path):
response = requests.get(url)
if response.status_code == 200:
with open(save_path, "wb") as file:
file.write(response.content)
return True
return False
def handler(event):
try:
image_url = event["input"]["image_url"]
# Create a temporary directory
os.makedirs("temp_images", exist_ok=True)
image_path = "temp_images/downloaded_image.jpg"
# Download the image
if not download_image(image_url, image_path):
raise Exception("Failed to download image")
# Process the image (your code here)
result = f"Processed image from: {image_url}"
# Cleanup specific folders after processing
clean(folder_list=["temp_images"])
return {"output": result}
except Exception as e:
# Attempt cleanup even if an error occurs
clean(folder_list=["temp_images"])
return {"error": str(e)}
runpod.serverless.start({"handler": handler})