import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import swin_t
from torchvision import transforms
from PIL import Image
import time

# Define MMIM model structure
class MMIM(nn.Module):
    def __init__(self, num_classes):
        super(MMIM, self).__init__()
        self.backbone = swin_t(weights=None)  # Load without pretrained weights
        self.backbone.head = nn.Identity()
        self.classifier = nn.Sequential(
            nn.Linear(768, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        x = self.backbone(x)
        return self.classifier(x)

# Load all models and compute class offsets
def load_all_models():
    model_defs = [
        ("MMIM_best1.pth", 9),   # Model 1
        ("MMIM_best3.pth", 4),   # Model 2
        ("MMIM_best2.pth", 12)   # Model 3
    ]
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    models = []
    offsets = []
    total_classes = 0

    for path, num_classes in model_defs:
        model = MMIM(num_classes)
        state_dict = torch.load(path, map_location=device)
        model.load_state_dict(state_dict)
        model.to(device)
        model.eval()
        models.append(model)
        offsets.append(total_classes)
        total_classes += num_classes

    idx_to_class = {i: f"class{i}" for i in range(total_classes)}
    return models, offsets, idx_to_class

# Predict image using confidence-based model router
def predict_image(image, models, offsets, idx_to_class, threshold=0.9):
    start = time.time()

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    transform = transforms.Compose([
        transforms.Resize((128, 128)),  # Faster inference
        transforms.ToTensor(),
        transforms.Normalize([0.5]*3, [0.5]*3)
    ])
    image_tensor = transform(i
runtime error
Exit code: 1. Reason: .5.4.mlp.0.bias", "backbone.features.5.4.mlp.3.weight", "backbone.features.5.4.mlp.3.bias", "backbone.features.5.5.norm1.weight", "backbone.features.5.5.norm1.bias", "backbone.features.5.5.attn.relative_position_bias_table", "backbone.features.5.5.attn.relative_position_index", "backbone.features.5.5.attn.qkv.weight", "backbone.features.5.5.attn.qkv.bias", "backbone.features.5.5.attn.proj.weight", "backbone.features.5.5.attn.proj.bias", "backbone.features.5.5.norm2.weight", "backbone.features.5.5.norm2.bias", "backbone.features.5.5.mlp.0.weight", "backbone.features.5.5.mlp.0.bias", "backbone.features.5.5.mlp.3.weight", "backbone.features.5.5.mlp.3.bias", "backbone.features.6.reduction.weight", "backbone.features.6.norm.weight", "backbone.features.6.norm.bias", "backbone.features.7.0.norm1.weight", "backbone.features.7.0.norm1.bias", "backbone.features.7.0.attn.relative_position_bias_table", "backbone.features.7.0.attn.relative_position_index", "backbone.features.7.0.attn.qkv.weight", "backbone.features.7.0.attn.qkv.bias", "backbone.features.7.0.attn.proj.weight", "backbone.features.7.0.attn.proj.bias", "backbone.features.7.0.norm2.weight", "backbone.features.7.0.norm2.bias", "backbone.features.7.0.mlp.0.weight", "backbone.features.7.0.mlp.0.bias", "backbone.features.7.0.mlp.3.weight", "backbone.features.7.0.mlp.3.bias", "backbone.features.7.1.norm1.weight", "backbone.features.7.1.norm1.bias", "backbone.features.7.1.attn.relative_position_bias_table", "backbone.features.7.1.attn.relative_position_index", "backbone.features.7.1.attn.qkv.weight", "backbone.features.7.1.attn.qkv.bias", "backbone.features.7.1.attn.proj.weight", "backbone.features.7.1.attn.proj.bias", "backbone.features.7.1.norm2.weight", "backbone.features.7.1.norm2.bias", "backbone.features.7.1.mlp.0.weight", "backbone.features.7.1.mlp.0.bias", "backbone.features.7.1.mlp.3.weight", "backbone.features.7.1.mlp.3.bias", "backbone.norm.weight", "backbone.norm.bias", "classifier.0.weight", "classifier.0.bias", "classifier.3.weight", "classifier.3.bias". 

Container logs:

===== Application Startup at 2025-06-26 06:03:02 =====

2025-06-26 08:04:04.986 WARNING streamlit.runtime.scriptrunner_utils.script_run_context: Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
2025-06-26 08:04:04.987 WARNING streamlit.runtime.scriptrunner_utils.script_run_context: Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
2025-06-26 08:04:04.999 
  Warning: to view this Streamlit app on a browser, run it with the following
  command:

    streamlit run app.py [ARGUMENTS]
2025-06-26 08:04:04.999 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
2025-06-26 08:04:04.999 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
2025-06-26 08:04:04.999 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
2025-06-26 08:04:04.999 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
2025-06-26 08:04:04.999 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
2025-06-26 08:04:05.000 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
2025-06-26 08:04:05.000 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
2025-06-26 08:04:05.000 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
2025-06-26 08:04:05.000 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
2025-06-26 08:04:05.548 Thread 'Thread-1': missing ScriptRunContext! This warning can be ignored when running in bare mode.
2025-06-26 08:04:05.548 Thread 'Thread-1': missing ScriptRunContext! This warning can be ignored when running in bare mode.
2025-06-26 08:04:05.548 Thread 'Thread-1': missing ScriptRunContext! This warning can be ignored when running in bare mode.
2025-06-26 08:04:06.139 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
2025-06-26 08:04:06.139 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
2025-06-26 08:04:06.139 Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
Traceback (most recent call last):
  File "/home/user/app/app.py", line 35, in <module>
    model_defs = load_all_models()
  File "/usr/local/lib/python3.10/site-packages/streamlit/runtime/caching/cache_utils.py", line 219, in __call__
    return self._get_or_create_cached_value(args, kwargs, spinner_message)
  File "/usr/local/lib/python3.10/site-packages/streamlit/runtime/caching/cache_utils.py", line 261, in _get_or_create_cached_value
    return self._handle_cache_miss(cache, value_key, func_args, func_kwargs)
  File "/usr/local/lib/python3.10/site-packages/streamlit/runtime/caching/cache_utils.py", line 320, in _handle_cache_miss
    computed_value = self._info.func(*func_args, **func_kwargs)
  File "/home/user/app/app.py", line 32, in load_all_models
    model_def["model"] = load_model(model_def["path"], len(model_def["class_names"]))
  File "/home/user/app/utils.py", line 9, in load_model
    model.load_state_dict(torch.load(model_path, map_location='cpu'))
  File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2593, in load_state_dict
    raise RuntimeError(
RuntimeError: Error(s) in loading state_dict for SwinTransformer:
	Missing key(s) in state_dict: "patch_embed.proj.weight", "patch_embed.proj.bias", "patch_embed.norm.weight", "patch_embed.norm.bias", "layers.0.blocks.0.norm1.weight", "layers.0.blocks.0.norm1.bias", "layers.0.blocks.0.attn.relative_position_bias_table", "layers.0.blocks.0.attn.qkv.weight", "layers.0.blocks.0.attn.qkv.bias", "layers.0.blocks.0.attn.proj.weight", "layers.0.blocks.0.attn.proj.bias", "layers.0.blocks.0.norm2.weight", "layers.0.blocks.0.norm2.bias", "layers.0.blocks.0.mlp.fc1.weight", "layers.0.blocks.0.mlp.fc1.bias", "layers.0.blocks.0.mlp.fc2.weight", "layers.0.blocks.0.mlp.fc2.bias", "layers.0.blocks.1.norm1.weight", "layers.0.blocks.1.norm1.bias", "layers.0.blocks.1.attn.relative_position_bias_table", "layers.0.blocks.1.attn.qkv.weight", "layers.0.blocks.1.attn.qkv.bias", "layers.0.blocks.1.attn.proj.weight", "layers.0.blocks.1.attn.proj.bias", "layers.0.blocks.1.norm2.weight", "layers.0.blocks.1.norm2.bias", "layers.0.blocks.1.mlp.fc1.weight", "layers.0.blocks.1.mlp.fc1.bias", "layers.0.blocks.1.mlp.fc2.weight", "layers.0.blocks.1.mlp.fc2.bias", "layers.1.downsample.norm.weight", "layers.1.downsample.norm.bias", "layers.1.downsample.reduction.weight", "layers.1.blocks.0.norm1.weight", "layers.1.blocks.0.norm1.bias", "layers.1.blocks.0.attn.relative_position_bias_table", "layers.1.blocks.0.attn.qkv.weight", "layers.1.blocks.0.attn.qkv.bias", "layers.1.blocks.0.attn.proj.weight", "layers.1.blocks.0.attn.proj.bias", "layers.1.blocks.0.norm2.weight", "layers.1.blocks.0.norm2.bias", "layers.1.blocks.0.mlp.fc1.weight", "layers.1.blocks.0.mlp.fc1.bias", "layers.1.blocks.0.mlp.fc2.weight", "layers.1.blocks.0.mlp.fc2.bias", "layers.1.blocks.1.norm1.weight", "layers.1.blocks.1.norm1.bias", "layers.1.blocks.1.attn.relative_position_bias_table", "layers.1.blocks.1.attn.qkv.weight", "layers.1.blocks.1.attn.qkv.bias", "layers.1.blocks.1.attn.proj.weight", "layers.1.blocks.1.attn.proj.bias", "layers.1.blocks.1.norm2.weight", "layers.1.blocks.1.norm2.bias", "layers.1.blocks.1.mlp.fc1.weight", "layers.1.blocks.1.mlp.fc1.bias", "layers.1.blocks.1.mlp.fc2.weight", "layers.1.blocks.1.mlp.fc2.bias", "layers.2.downsample.norm.weight", "layers.2.downsample.norm.bias", "layers.2.downsample.reduction.weight", "layers.2.blocks.0.norm1.weight", "layers.2.blocks.0.norm1.bias", "layers.2.blocks.0.attn.relative_position_bias_table", "layers.2.blocks.0.attn.qkv.weight", "layers.2.blocks.0.attn.qkv.bias", "layers.2.blocks.0.attn.proj.weight", "layers.2.blocks.0.attn.proj.bias", "layers.2.blocks.0.norm2.weight", "layers.2.blocks.0.norm2.bias", "layers.2.blocks.0.mlp.fc1.weight", "layers.2.blocks.0.mlp.fc1.bias", "layers.2.blocks.0.mlp.fc2.weight", "layers.2.blocks.0.mlp.fc2.bias", "layers.2.blocks.1.norm1.weight", "layers.2.blocks.1.norm1.bias", "layers.2.blocks.1.attn.relative_position_bias_table", "layers.2.blocks.1.attn.qkv.weight", "layers.2.blocks.1.attn.qkv.bias", "layers.2.blocks.1.attn.proj.weight", "layers.2.blocks.1.attn.proj.bias", "layers.2.blocks.1.norm2.weight", "layers.2.blocks.1.norm2.bias", "layers.2.blocks.1.mlp.fc1.weight", "layers.2.blocks.1.mlp.fc1.bias", "layers.2.blocks.1.mlp.fc2.weight", "layers.2.blocks.1.mlp.fc2.bias", "layers.2.blocks.2.norm1.weight", "layers.2.blocks.2.norm1.bias", "layers.2.blocks.2.attn.relative_position_bias_table", "layers.2.blocks.2.attn.qkv.weight", "layers.2.blocks.2.attn.qkv.bias", "layers.2.blocks.2.attn.proj.weight", "layers.2.blocks.2.attn.proj.bias", "layers.2.blocks.2.norm2.weight", "layers.2.blocks.2.norm2.bias", "layers.2.blocks.2.mlp.fc1.weight", "layers.2.blocks.2.mlp.fc1.bias", "layers.2.blocks.2.mlp.fc2.weight", "layers.2.blocks.2.mlp.fc2.bias", "layers.2.blocks.3.norm1.weight", "layers.2.blocks.3.norm1.bias", "layers.2.blocks.3.attn.relative_position_bias_table", "layers.2.blocks.3.attn.qkv.weight", "layers.2.blocks.3.attn.qkv.bias", "layers.2.blocks.3.attn.proj.weight", "layers.2.blocks.3.attn.proj.bias", "layers.2.blocks.3.norm2.weight", "layers.2.blocks.3.norm2.bias", "layers.2.blocks.3.mlp.fc1.weight", "layers.2.blocks.3.mlp.fc1.bias", "layers.2.blocks.3.mlp.fc2.weight", "layers.2.blocks.3.mlp.fc2.bias", "layers.2.blocks.4.norm1.weight", "layers.2.blocks.4.norm1.bias", "layers.2.blocks.4.attn.relative_position_bias_table", "layers.2.blocks.4.attn.qkv.weight", "layers.2.blocks.4.attn.qkv.bias", "layers.2.blocks.4.attn.proj.weight", "layers.2.blocks.4.attn.proj.bias", "layers.2.blocks.4.norm2.weight", "layers.2.blocks.4.norm2.bias", "layers.2.blocks.4.mlp.fc1.weight", "layers.2.blocks.4.mlp.fc1.bias", "layers.2.blocks.4.mlp.fc2.weight", "layers.2.blocks.4.mlp.fc2.bias", "layers.2.blocks.5.norm1.weight", "layers.2.blocks.5.norm1.bias", "layers.2.blocks.5.attn.relative_position_bias_table", "layers.2.blocks.5.attn.qkv.weight", "layers.2.blocks.5.attn.qkv.bias", "layers.2.blocks.5.attn.proj.weight", "layers.2.blocks.5.attn.proj.bias", "layers.2.blocks.5.norm2.weight", "layers.2.blocks.5.norm2.bias", "layers.2.blocks.5.mlp.fc1.weight", "layers.2.blocks.5.mlp.fc1.bias", "layers.2.blocks.5.mlp.fc2.weight", "layers.2.blocks.5.mlp.fc2.bias", "layers.3.downsample.norm.weight", "layers.3.downsample.norm.bias", "layers.3.downsample.reduction.weight", "layers.3.blocks.0.norm1.weight", "layers.3.blocks.0.norm1.bias", "layers.3.blocks.0.attn.relative_position_bias_table", "layers.3.blocks.0.attn.qkv.weight", "layers.3.blocks.0.attn.qkv.bias", "layers.3.blocks.0.attn.proj.weight", "layers.3.blocks.0.attn.proj.bias", "layers.3.blocks.0.norm2.weight", "layers.3.blocks.0.norm2.bias", "layers.3.blocks.0.mlp.fc1.weight", "layers.3.blocks.0.mlp.fc1.bias", "layers.3.blocks.0.mlp.fc2.weight", "layers.3.blocks.0.mlp.fc2.bias", "layers.3.blocks.1.norm1.weight", "layers.3.blocks.1.norm1.bias", "layers.3.blocks.1.attn.relative_position_bias_table", "layers.3.blocks.1.attn.qkv.weight", "layers.3.blocks.1.attn.qkv.bias", "layers.3.blocks.1.attn.proj.weight", "layers.3.blocks.1.attn.proj.bias", "layers.3.blocks.1.norm2.weight", "layers.3.blocks.1.norm2.bias", "layers.3.blocks.1.mlp.fc1.weight", "layers.3.blocks.1.mlp.fc1.bias", "layers.3.blocks.1.mlp.fc2.weight", "layers.3.blocks.1.mlp.fc2.bias", "norm.weight", "norm.bias", "head.fc.weight", "head.fc.bias". 
	Unexpected key(s) in state_dict: "backbone.features.0.0.weight", "backbone.features.0.0.bias", "backbone.features.0.2.weight", "backbone.features.0.2.bias", "backbone.features.1.0.norm1.weight", "backbone.features.1.0.norm1.bias", "backbone.features.1.0.attn.relative_position_bias_table", "backbone.features.1.0.attn.relative_position_index", "backbone.features.1.0.attn.qkv.weight", "backbone.features.1.0.attn.qkv.bias", "backbone.features.1.0.attn.proj.weight", "backbone.features.1.0.attn.proj.bias", "backbone.features.1.0.norm2.weight", "backbone.features.1.0.norm2.bias", "backbone.features.1.0.mlp.0.weight", "backbone.features.1.0.mlp.0.bias", "backbone.features.1.0.mlp.3.weight", "backbone.features.1.0.mlp.3.bias", "backbone.features.1.1.norm1.weight", "backbone.features.1.1.norm1.bias", "backbone.features.1.1.attn.relative_position_bias_table", "backbone.features.1.1.attn.relative_position_index", "backbone.features.1.1.attn.qkv.weight", "backbone.features.1.1.attn.qkv.bias", "backbone.features.1.1.attn.proj.weight", "backbone.features.1.1.attn.proj.bias", "backbone.features.1.1.norm2.weight", "backbone.features.1.1.norm2.bias", "backbone.features.1.1.mlp.0.weight", "backbone.features.1.1.mlp.0.bias", "backbone.features.1.1.mlp.3.weight", "backbone.features.1.1.mlp.3.bias", "backbone.features.2.reduction.weight", "backbone.features.2.norm.weight", "backbone.features.2.norm.bias", "backbone.features.3.0.norm1.weight", "backbone.features.3.0.norm1.bias", "backbone.features.3.0.attn.relative_position_bias_table", "backbone.features.3.0.attn.relative_position_index", "backbone.features.3.0.attn.qkv.weight", "backbone.features.3.0.attn.qkv.bias", "backbone.features.3.0.attn.proj.weight", "backbone.features.3.0.attn.proj.bias", "backbone.features.3.0.norm2.weight", "backbone.features.3.0.norm2.bias", "backbone.features.3.0.mlp.0.weight", "backbone.features.3.0.mlp.0.bias", "backbone.features.3.0.mlp.3.weight", "backbone.features.3.0.mlp.3.bias", "backbone.features.3.1.norm1.weight", "backbone.features.3.1.norm1.bias", "backbone.features.3.1.attn.relative_position_bias_table", "backbone.features.3.1.attn.relative_position_index", "backbone.features.3.1.attn.qkv.weight", "backbone.features.3.1.attn.qkv.bias", "backbone.features.3.1.attn.proj.weight", "backbone.features.3.1.attn.proj.bias", "backbone.features.3.1.norm2.weight", "backbone.features.3.1.norm2.bias", "backbone.features.3.1.mlp.0.weight", "backbone.features.3.1.mlp.0.bias", "backbone.features.3.1.mlp.3.weight", "backbone.features.3.1.mlp.3.bias", "backbone.features.4.reduction.weight", "backbone.features.4.norm.weight", "backbone.features.4.norm.bias", "backbone.features.5.0.norm1.weight", "backbone.features.5.0.norm1.bias", "backbone.features.5.0.attn.relative_position_bias_table", "backbone.features.5.0.attn.relative_position_index", "backbone.features.5.0.attn.qkv.weight", "backbone.features.5.0.attn.qkv.bias", "backbone.features.5.0.attn.proj.weight", "backbone.features.5.0.attn.proj.bias", "backbone.features.5.0.norm2.weight", "backbone.features.5.0.norm2.bias", "backbone.features.5.0.mlp.0.weight", "backbone.features.5.0.mlp.0.bias", "backbone.features.5.0.mlp.3.weight", "backbone.features.5.0.mlp.3.bias", "backbone.features.5.1.norm1.weight", "backbone.features.5.1.norm1.bias", "backbone.features.5.1.attn.relative_position_bias_table", "backbone.features.5.1.attn.relative_position_index", "backbone.features.5.1.attn.qkv.weight", "backbone.features.5.1.attn.qkv.bias", "backbone.features.5.1.attn.proj.weight", "backbone.features.5.1.attn.proj.bias", "backbone.features.5.1.norm2.weight", "backbone.features.5.1.norm2.bias", "backbone.features.5.1.mlp.0.weight", "backbone.features.5.1.mlp.0.bias", "backbone.features.5.1.mlp.3.weight", "backbone.features.5.1.mlp.3.bias", "backbone.features.5.2.norm1.weight", "backbone.features.5.2.norm1.bias", "backbone.features.5.2.attn.relative_position_bias_table", "backbone.features.5.2.attn.relative_position_index", "backbone.features.5.2.attn.qkv.weight", "backbone.features.5.2.attn.qkv.bias", "backbone.features.5.2.attn.proj.weight", "backbone.features.5.2.attn.proj.bias", "backbone.features.5.2.norm2.weight", "backbone.features.5.2.norm2.bias", "backbone.features.5.2.mlp.0.weight", "backbone.features.5.2.mlp.0.bias", "backbone.features.5.2.mlp.3.weight", "backbone.features.5.2.mlp.3.bias", "backbone.features.5.3.norm1.weight", "backbone.features.5.3.norm1.bias", "backbone.features.5.3.attn.relative_position_bias_table", "backbone.features.5.3.attn.relative_position_index", "backbone.features.5.3.attn.qkv.weight", "backbone.features.5.3.attn.qkv.bias", "backbone.features.5.3.attn.proj.weight", "backbone.features.5.3.attn.proj.bias", "backbone.features.5.3.norm2.weight", "backbone.features.5.3.norm2.bias", "backbone.features.5.3.mlp.0.weight", "backbone.features.5.3.mlp.0.bias", "backbone.features.5.3.mlp.3.weight", "backbone.features.5.3.mlp.3.bias", "backbone.features.5.4.norm1.weight", "backbone.features.5.4.norm1.bias", "backbone.features.5.4.attn.relative_position_bias_table", "backbone.features.5.4.attn.relative_position_index", "backbone.features.5.4.attn.qkv.weight", "backbone.features.5.4.attn.qkv.bias", "backbone.features.5.4.attn.proj.weight", "backbone.features.5.4.attn.proj.bias", "backbone.features.5.4.norm2.weight", "backbone.features.5.4.norm2.bias", "backbone.features.5.4.mlp.0.weight", "backbone.features.5.4.mlp.0.bias", "backbone.features.5.4.mlp.3.weight", "backbone.features.5.4.mlp.3.bias", "backbone.features.5.5.norm1.weight", "backbone.features.5.5.norm1.bias", "backbone.features.5.5.attn.relative_position_bias_table", "backbone.features.5.5.attn.relative_position_index", "backbone.features.5.5.attn.qkv.weight", "backbone.features.5.5.attn.qkv.bias", "backbone.features.5.5.attn.proj.weight", "backbone.features.5.5.attn.proj.bias", "backbone.features.5.5.norm2.weight", "backbone.features.5.5.norm2.bias", "backbone.features.5.5.mlp.0.weight", "backbone.features.5.5.mlp.0.bias", "backbone.features.5.5.mlp.3.weight", "backbone.features.5.5.mlp.3.bias", "backbone.features.6.reduction.weight", "backbone.features.6.norm.weight", "backbone.features.6.norm.bias", "backbone.features.7.0.norm1.weight", "backbone.features.7.0.norm1.bias", "backbone.features.7.0.attn.relative_position_bias_table", "backbone.features.7.0.attn.relative_position_index", "backbone.features.7.0.attn.qkv.weight", "backbone.features.7.0.attn.qkv.bias", "backbone.features.7.0.attn.proj.weight", "backbone.features.7.0.attn.proj.bias", "backbone.features.7.0.norm2.weight", "backbone.features.7.0.norm2.bias", "backbone.features.7.0.mlp.0.weight", "backbone.features.7.0.mlp.0.bias", "backbone.features.7.0.mlp.3.weight", "backbone.features.7.0.mlp.3.bias", "backbone.features.7.1.norm1.weight", "backbone.features.7.1.norm1.bias", "backbone.features.7.1.attn.relative_position_bias_table", "backbone.features.7.1.attn.relative_position_index", "backbone.features.7.1.attn.qkv.weight", "backbone.features.7.1.attn.qkv.bias", "backbone.features.7.1.attn.proj.weight", "backbone.features.7.1.attn.proj.bias", "backbone.features.7.1.norm2.weight", "backbone.features.7.1.norm2.bias", "backbone.features.7.1.mlp.0.weight", "backbone.features.7.1.mlp.0.bias", "backbone.features.7.1.mlp.3.weight", "backbone.features.7.1.mlp.3.bias", "backbone.norm.weight", "backbone.norm.bias", "classifier.0.weight", "classifier.0.bias", "classifier.3.weight", "classifier.3.bias". 

mage).unsqueeze(0).to(device)

    temperatures = [1.2, 0.9, 1.0]  # Adjust if needed

    max_score = float('-inf')
    final_pred = -1
    confidence = 0.0

    for model, offset, temp in zip(models, offsets, temperatures):
        with torch.no_grad():
            logits = model(image_tensor) / temp
            probs = F.softmax(logits, dim=1).squeeze(0)
            score, pred = torch.max(probs, dim=0)
            if score.item() > max_score:
                max_score = score.item()
                final_pred = pred.item() + offset
                confidence = score.item()

    print(f"Inference took: {time.time() - start:.2f} seconds")  # Log time

    if confidence >= threshold:
        return {idx_to_class[final_pred]: round(confidence, 4)}
    else:
        return {"Low confidence": round(confidence, 4)}