Views
No views yet
hepg2 - HepG2 cell linek562 - K562 cell linewtc11 - WTC11 cell line1from model_loader import load_cell_type_model
2
3# Load model for HepG2
4model = load_cell_type_model("hepg2")
5
6# Load model for K562
7model = load_cell_type_model("k562")1def get_device():
2 """Automatically detects available device"""
3 if torch.cuda.is_available():
4 return torch.device("cuda")
5 else:
6 return torch.device("cpu")
7
8# Load Pre-Trained Model Weights for Human Legnet
9def download_and_load_model(cell_type="k562", repo_id="Ni-os/MPRALegNet", device=None):
10 # Download main config
11 config_path = hf_hub_download(
12 repo_id=repo_id,
13 filename="config.json"
14 )
15
16 # Load config
17 with open(config_path, 'r') as f:
18 config = json.load(f)
19
20 # Create model
21 model = LegNet(
22 in_ch=config["in_ch"],
23 stem_ch=config["stem_ch"],
24 stem_ks=config["stem_ks"],
25 ef_ks=config["ef_ks"],
26 ef_block_sizes=config["ef_block_sizes"],
27 pool_sizes=config["pool_sizes"],
28 resize_factor=config["resize_factor"],
29 activation=torch.nn.SiLU
30 ).to(device)
31
32 # Determine which weight file to download
33 weight_files = {
34 "hepg2": "weights/hepg2_best_model_test1_val2.safetensors",
35 "k562": "weights/k562_best_model_test1_val2.safetensors",
36 "wtc11": "weights/wtc11_best_model_test1_val2.safetensors"
37 }
38
39 # Download weights
40 weights_path = hf_hub_download(
41 repo_id=repo_id,
42 filename=weight_files[cell_type.lower()]
43 )
44
45 # Load weights into model
46 state_dict = load_file(weights_path)
47 model.load_state_dict(state_dict)
48 model.eval()
49 print(f"✅ Model for {cell_type} loaded!")
50 return model
51
52device = get_device()
53
54print("Loading pre-trained model weights for Human Legnet")
55model_legnet = download_and_load_model("hepg2", device = device)