Update app.py

app.py

@@ -54,7 +54,7 @@ def load_custom_model():
 def process_p_file(uploaded_file, percentage_idx, complexity_idx):
     capture = PrintCapture()
     sys.stdout = capture  # Redirect print statements to capture
-    
+
     try:
         model_repo_url = "https://huggingface.co/sadjadalikhani/LWM"
         model_repo_dir = "./LWM"
@@ -69,9 +69,11 @@ def process_p_file(uploaded_file, percentage_idx, complexity_idx):
         else:
             return f"Directory {model_repo_dir} does not exist."
 
-        print("a")
-        from lwm_model import LWM
-        print("b")
+        # Add LWM repo path to Python module search path
+        if model_repo_dir not in sys.path:
+            sys.path.append(model_repo_dir)
+
+        from lwm_model import LWM  # Now this should work
         device = 'cpu'
         print(f"Loading the LWM model on {device}...")
         model = LWM.from_pretrained(device=device)