Update app.py

app.py

@@ -34,20 +34,45 @@ from transformers import AutoModel  # Assuming you use a transformer-like model
 import numpy as np
 import importlib.util
 
-# Function to load the pre-trained model from Hugging Face
-def load_pretrained_model():
-    # Load the pre-trained model from the Hugging Face repo
-    model = AutoModel.from_pretrained("sadjadalikhani/LWM")
-    model.eval()  # Set model to evaluation mode
+import torch
+import numpy as np
+import importlib.util
+import subprocess
+import os
+
+# Function to load the pre-trained model from your cloned repository
+def load_custom_model():
+    # Assume your model is in the cloned LWM repository
+    from lwm_model import LWM  # Assuming the model is defined in lwm_model.py
+    model = LWM()  # Modify this according to your model initialization
+    model.eval()  # Set the model to evaluation mode
     return model
 
-# Function to process the uploaded .py file and perform inference using the model
+# Function to process the uploaded .py file and perform inference using the custom model
 def process_python_file(uploaded_file, percentage_idx, complexity_idx):
     try:
-        # Step 1: Load the model
-        model = load_pretrained_model()
+        # Clone the repository if not already done (for model and tokenizer)
+        model_repo_url = "https://huggingface.co/sadjadalikhani/LWM"
+        model_repo_dir = "./LWM"
+
+        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)
 
-        # Step 2: Load the uploaded .py file that contains the wireless channel matrix
+        # Change the working directory to the cloned LWM folder
+        if os.path.exists(model_repo_dir):
+            os.chdir(model_repo_dir)
+            print(f"Changed working directory to {os.getcwd()}")
+        else:
+            return f"Directory {model_repo_dir} does not exist."
+
+        # Step 1: Load the custom model
+        model = load_custom_model()
+
+        # Step 2: Import the tokenizer
+        from input_preprocess import tokenizer
+
+        # Step 3: Load the uploaded .py file that contains the wireless channel matrix
         # Import the Python file dynamically
         spec = importlib.util.spec_from_file_location("uploaded_module", uploaded_file.name)
         uploaded_module = importlib.util.module_from_spec(spec)
@@ -56,13 +81,15 @@ def process_python_file(uploaded_file, percentage_idx, complexity_idx):
         # Assuming the uploaded file defines a variable called 'channel_matrix'
         channel_matrix = uploaded_module.channel_matrix  # This should be defined in the uploaded file
 
-        # Step 3: Perform inference on the channel matrix using the model
+        # Step 4: Tokenize the data if needed (or perform any necessary preprocessing)
+        preprocessed_data = tokenizer(manual_data=channel_matrix, gen_raw=True)
+
+        # Step 5: Perform inference on the channel matrix using the model
         with torch.no_grad():
-            input_tensor = torch.tensor(channel_matrix).unsqueeze(0)  # Add batch dimension
+            input_tensor = torch.tensor(preprocessed_data).unsqueeze(0)  # Add batch dimension
             output = model(input_tensor)  # Perform inference
 
-        # Step 4: Generate new images based on the inference results
-        # You can modify this logic depending on how you want to visualize the results
+        # Step 6: Generate new images based on the inference results
         generated_raw_img = np.random.rand(300, 300, 3) * 255  # Placeholder: Replace with actual inference result
         generated_embeddings_img = np.random.rand(300, 300, 3) * 255  # Placeholder: Replace with actual inference result
         
@@ -82,7 +109,6 @@ def process_python_file(uploaded_file, percentage_idx, complexity_idx):
     except Exception as e:
         return str(e), str(e)
 
-
 # Function to handle logic based on whether a file is uploaded or not
 def los_nlos_classification(file, percentage_idx, complexity_idx):
     if file is not None: