import torch
import torch.nn as nn
import torch.nn.functional as F

# Simplified U2Net architecture would go here
# For actual implementation, you'd want to import from the original repo

def load_model(model_name):
    model_urls = {
        "u2net": "https://github.com/xuebinqin/U-2-Net/raw/master/models/u2net.pth",
        "u2net_human_seg": "https://github.com/xuebinqin/U-2-Net/raw/master/models/u2net_human_seg.pth"
    }
    model = U2Net()  # This would be the actual U2Net class
    model.load_state_dict(torch.hub.load_state_dict_from_url(model_urls[model_name]))
    return model