Update app.py

app.py

@@ -30,7 +30,6 @@ def display_predefined_images(percentage_idx, complexity_idx):
     return raw_image, embeddings_image
 
 import torch
-from transformers import AutoModel  # Assuming you use a transformer-like model in your LWM repo
 import numpy as np
 import importlib.util
 
@@ -67,44 +66,57 @@ def process_python_file(uploaded_file, percentage_idx, complexity_idx):
             return f"Directory {model_repo_dir} does not exist."
 
         # Step 1: Load the custom model
-        model = load_custom_model()
+        from lwm_model import LWM
+        import torch
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        print(f"Loading the LWM model on {device}...")
+        model = LWM.from_pretrained(device=device)
 
         # 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)
-        spec.loader.exec_module(uploaded_module)
+        ## 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)
+        #spec.loader.exec_module(uploaded_module)
 
-        # Assuming the uploaded file defines a variable called 'channel_matrix'
-        channel_matrix = uploaded_module.channel_matrix  # This should be defined in the uploaded file
+        ## Assuming the uploaded file defines a variable called 'dataset'
+        #manual_data = uploaded_module.dataset  # This should be defined in the uploaded file
+
+        import pickle
+        # Load the .p file containing the wireless channel matrix
+        with open(uploaded_file.name, 'rb') as f:
+            manual_data = pickle.load(f)
 
         # Step 4: Tokenize the data if needed (or perform any necessary preprocessing)
-        preprocessed_data = tokenizer(manual_data=channel_matrix, gen_raw=True)
+        preprocessed_chs = tokenizer(manual_data=manual_data)
 
         # Step 5: Perform inference on the channel matrix using the model
-        with torch.no_grad():
-            input_tensor = torch.tensor(preprocessed_data).unsqueeze(0)  # Add batch dimension
-            output = model(input_tensor)  # Perform inference
-
+        #with torch.no_grad():
+        #    input_tensor = torch.tensor(preprocessed_data).unsqueeze(0)  # Add batch dimension
+        #    output = model(input_tensor)  # Perform inference
+        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)
+        print(output_emb.shape)
+        print(output_raw.shape)
         # 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
+        #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
         
         # Save the generated images
-        generated_raw_image_path = os.path.join(GENERATED_PATH, f"generated_raw_{percentage_idx}_{complexity_idx}.png")
-        generated_embeddings_image_path = os.path.join(GENERATED_PATH, f"generated_embeddings_{percentage_idx}_{complexity_idx}.png")
+        #generated_raw_image_path = os.path.join(GENERATED_PATH, f"generated_raw_{percentage_idx}_{complexity_idx}.png")
+        #generated_embeddings_image_path = os.path.join(GENERATED_PATH, f"generated_embeddings_{percentage_idx}_{complexity_idx}.png")
         
-        Image.fromarray(generated_raw_img.astype(np.uint8)).save(generated_raw_image_path)
-        Image.fromarray(generated_embeddings_img.astype(np.uint8)).save(generated_embeddings_image_path)
+        #Image.fromarray(generated_raw_img.astype(np.uint8)).save(generated_raw_image_path)
+        #Image.fromarray(generated_embeddings_img.astype(np.uint8)).save(generated_embeddings_image_path)
 
         # Load the generated images
-        raw_image = Image.open(generated_raw_image_path)
-        embeddings_image = Image.open(generated_embeddings_image_path)
+        #raw_image = Image.open(generated_raw_image_path)
+        #embeddings_image = Image.open(generated_embeddings_image_path)
 
-        return raw_image, embeddings_image
+        return output_emb, output_raw
 
     except Exception as e:
         return str(e), str(e)