Update handler.py

handler.py

@@ -6,15 +6,32 @@ from io import BytesIO
 from PIL import Image
 import requests
 from transformers import AutoModelForCausalLM, AutoProcessor
-def install(package):
-    subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-warn-script-location", package])
+import os
+
+def create_venv(venv_path="env"):
+    """Creates a virtual environment."""
+    if not os.path.exists(venv_path):
+        print(f"[INFO] Creating virtual environment at {venv_path}")
+        subprocess.check_call([sys.executable, "-m", "venv", venv_path])
+    else:
+        print(f"[INFO] Virtual environment already exists at {venv_path}")
+
+def install(package, venv_path="env"):
+    """Installs a package inside the virtual environment."""
+    pip_path = os.path.join(venv_path, "bin", "pip")
+    subprocess.check_call([pip_path, "install", "--no-warn-script-location", package])
 
 class EndpointHandler:
     def __init__(self, path=""):
+        # Step 1: Create the virtual environment
+        self.venv_path = "env"
+        create_venv(self.venv_path)
+
+        # Step 2: Install required packages inside the virtual environment
         required_packages = ['timm', 'einops', 'flash-attn', 'Pillow']
         for package in required_packages:
             try:
-                install(package)
+                install(package, self.venv_path)
                 print(f"Successfully installed {package}")
             except Exception as e:
                 print(f"Failed to install {package}: {str(e)}")