Update app.py

app.py

@@ -54,63 +54,35 @@ 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"
 
-        # Step 1: Clone the model repository if not already cloned
         if not os.path.exists(model_repo_dir):
             print(f"Cloning model repository from {model_repo_url}...")
             subprocess.run(["git", "clone", model_repo_url, model_repo_dir], check=True)
-        
-        # Debugging: Check if the directory exists and print contents
+
         if os.path.exists(model_repo_dir):
             os.chdir(model_repo_dir)
             print(f"Changed working directory to {os.getcwd()}")
-            print(f"Directory content: {os.listdir(os.getcwd())}")  # Debugging: Check repo content
         else:
-            print(f"Directory {model_repo_dir} does not exist.")
-            return
+            return f"Directory {model_repo_dir} does not exist."
 
-        # Step 2: Add the cloned repo to sys.path for imports
-        if model_repo_dir not in sys.path:
-            sys.path.append(model_repo_dir)
-        
-        # Debugging: Print sys.path to ensure the cloned repo is in the path
-        print(f"sys.path: {sys.path}")
-
-        # Step 3: Dynamically import the model after cloning
-        try:
-            from lwm_model import LWM  # Custom model in the cloned repo
-            print("Successfully imported LWM model.")
-        except ImportError as e:
-            print(f"Error importing LWM model: {e}")
-            print("Make sure lwm_model.py exists in the cloned repository.")
-            return
-
-        # Step 4: Check if GPU is available and set the device
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        print(f"Using device: {device}")
-
-        # Load the model from the cloned repository
+        print("a")
+        from lwm_model import LWM
+        print("b")
+        device = 'cpu'
+        print(f"Loading the LWM model on {device}...")
         model = LWM.from_pretrained(device=device)
 
-        # Step 5: Import the tokenizer
-        try:
-            from input_preprocess import tokenizer
-        except ImportError as e:
-            print(f"Error importing tokenizer: {e}")
-            return
+        from input_preprocess import tokenizer
 
-        # Step 6: Load the uploaded .p file (wireless channel matrix)
         with open(uploaded_file.name, 'rb') as f:
             manual_data = pickle.load(f)
 
-        # Step 7: Tokenize the data if needed (or perform any necessary preprocessing)
         preprocessed_chs = tokenizer(manual_data=manual_data)
 
-        # Step 8: Perform inference using the model
         from inference import lwm_inference, create_raw_dataset
         output_emb = lwm_inference(preprocessed_chs, 'channel_emb', model)
         output_raw = create_raw_dataset(preprocessed_chs, device)
@@ -118,15 +90,13 @@ def process_p_file(uploaded_file, percentage_idx, complexity_idx):
         print(f"Output Embeddings Shape: {output_emb.shape}")
         print(f"Output Raw Shape: {output_raw.shape}")
 
-        # Return the embeddings, raw output, and captured output
         return output_emb, output_raw, capture.get_output()
 
     except Exception as e:
-        # Handle exceptions and return the captured output
         return str(e), str(e), capture.get_output()
 
     finally:
-        sys.stdout = sys.__stdout__  # Reset stdout
+        sys.stdout = sys.__stdout__  # Reset print statements
 
 # Function to handle logic based on whether a file is uploaded or not
 def los_nlos_classification(file, percentage_idx, complexity_idx):