Update app.py

app.py

@@ -50,6 +50,13 @@ def load_custom_model():
     model.eval()
     return model
 
+import sys
+import subprocess
+import os
+import pickle
+import torch
+import io
+
 # Function to process the uploaded .p file and perform inference using the custom model
 def process_p_file(uploaded_file, percentage_idx, complexity_idx):
     capture = PrintCapture()
@@ -59,32 +66,58 @@ def process_p_file(uploaded_file, percentage_idx, complexity_idx):
         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:
-            return f"Directory {model_repo_dir} does not exist."
+            print(f"Directory {model_repo_dir} does not exist.")
+            return
 
-        # Add LWM repo path to Python module search path
+        # 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)
-
-        from lwm_model import LWM  # Now this should work
-        device = 'cpu'
-        print(f"Loading the LWM model on {device}...")
+        
+        # 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
         model = LWM.from_pretrained(device=device)
 
-        from input_preprocess import tokenizer
+        # Step 5: Import the tokenizer
+        try:
+            from input_preprocess import tokenizer
+        except ImportError as e:
+            print(f"Error importing tokenizer: {e}")
+            return
 
+        # 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)
@@ -92,13 +125,15 @@ 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 print statements
+        sys.stdout = sys.__stdout__  # Reset stdout
 
 # Function to handle logic based on whether a file is uploaded or not
 def los_nlos_classification(file, percentage_idx, complexity_idx):