Spaces:
Running on Zero
Running on Zero
| import torch | |
| from utils.fxencoder_plusplus import FxEncoderPlusPlus | |
| from utils.fxencoder_plusplus.model import get_model_path | |
| def load_model(model_name="default", model_path=None, device="cuda", auto_download=True, cache_dir=None): | |
| """ | |
| Load FxEncoderPlusPlus model. | |
| Args: | |
| model_name: Name of pretrained model ('default', 'musdb', 'medleydb') | |
| model_path: Custom checkpoint path. If provided, ignores model_name | |
| device: Device to load model on ('cuda' or 'cpu') | |
| auto_download: Automatically download if model not found | |
| cache_dir: Custom cache directory for downloaded models | |
| Returns: | |
| Loaded FxEncoderPlusPlus model | |
| Examples: | |
| # Load default base model | |
| model = load_model() | |
| # Load musdb model | |
| model = load_model(model_name="musdb") | |
| # Load medleydb model | |
| model = load_model(model_name="medleydb") | |
| # Load custom checkpoint | |
| model = load_model(model_path="/path/to/custom.pt") | |
| # List available models | |
| list_available_models() | |
| """ | |
| # Handle device | |
| if device == "cuda" and not torch.cuda.is_available(): | |
| print("CUDA not available, using CPU") | |
| device = "cpu" | |
| # Determine model path | |
| if model_path is None: | |
| if auto_download: | |
| model_path = get_model_path(model_name, cache_dir=cache_dir) | |
| else: | |
| raise ValueError("model_path is None and auto_download is False") | |
| # Create model instance with specified device | |
| model = FxEncoderPlusPlus( | |
| embed_dim=2048, | |
| audio_clap_module=False, | |
| text_clap_module=False, | |
| extractor_module=False, | |
| device=device | |
| ) | |
| # Load checkpoint | |
| checkpoint = torch.load(model_path, map_location=device, weights_only=False) | |
| if "epoch" in checkpoint: | |
| # resuming a train checkpoint w/ epoch and optimizer state | |
| start_epoch = checkpoint["epoch"] | |
| sd = checkpoint["state_dict"] | |
| if next(iter(sd.items()))[0].startswith("module"): | |
| sd = {k[len("module."):]: v for k, v in sd.items()} | |
| model.load_state_dict(sd, strict=False ) | |
| print(f"Loaded checkpoint from epoch {start_epoch}") | |
| else: | |
| # loading a bare (model only) checkpoint for fine-tune or evaluation | |
| model.load_state_dict(checkpoint) | |
| print("Loaded model checkpoint") | |
| model.to(device) | |
| model.eval() | |
| # Freeze parameters for inference | |
| for param in model.parameters(): | |
| param.requires_grad = False | |
| print(f"Model loaded successfully on {device}") | |
| return model | |