Spaces:
Running
Running
| import torch | |
| from transformers import AutoModel # Assuming you use a transformer-like model in your LWM repo | |
| 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 | |
| return model | |
| # Function to process the uploaded .py file and perform inference using the model | |
| def process_python_file(uploaded_file, percentage_idx, complexity_idx): | |
| try: | |
| # Step 1: Load the model | |
| model = load_pretrained_model() | |
| # Step 2: 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 | |
| # Step 3: Perform inference on the channel matrix using the model | |
| with torch.no_grad(): | |
| input_tensor = torch.tensor(channel_matrix).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 | |
| 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") | |
| 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) | |
| return raw_image, embeddings_image | |
| except Exception as e: | |
| return str(e), str(e) | |