From ef714f010e4f808226168c210f0c6a34f1552b49 Mon Sep 17 00:00:00 2001 From: Manan17 Date: Sun, 8 Mar 2026 04:18:20 +0000 Subject: [PATCH] adding export support --- studio/backend/core/export/export.py | 89 ++++++++++++++++++++++++---- 1 file changed, 77 insertions(+), 12 deletions(-) diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 3d3db2d560..2393b03e22 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -17,6 +17,7 @@ import torch from utils.hardware import clear_gpu_cache from utils.models import is_vision_model, get_base_model_from_lora +from utils.models.model_config import detect_audio_type from core.inference import get_inference_backend logger = logging.getLogger(__name__) @@ -96,6 +97,7 @@ class ExportBackend: self.current_tokenizer = None self.is_vision = False self.is_peft = False + self._audio_type = None def cleanup_memory(self): """Offload and delete all models from memory""" @@ -111,6 +113,7 @@ class ExportBackend: self.current_model = None self.current_tokenizer = None self.current_checkpoint = None + self._audio_type = None # Clear GPU memory cache (handles gc + backend-specific cleanup) clear_gpu_cache() @@ -148,24 +151,75 @@ class ExportBackend: # First, cleanup existing models self.cleanup_memory() - # Detect if vision model checkpoint_path_obj = Path(checkpoint_path) - # Check if it's a LoRA adapter + # Determine the model identity for type detection adapter_config = checkpoint_path_obj / "adapter_config.json" + base_model = None if adapter_config.exists(): - # It's a LoRA - get base model to check vision base_model = get_base_model_from_lora(checkpoint_path) - if base_model: - self.is_vision = is_vision_model(base_model) - else: + if not base_model: return False, "Could not determine base model for adapter" - else: - # Check the model itself - self.is_vision = is_vision_model(checkpoint_path) + + model_id = base_model or checkpoint_path + + # Detect audio type and vision + self._audio_type = detect_audio_type(model_id) + self.is_vision = not self._audio_type and is_vision_model(model_id) # Load model based on type - if self.is_vision: + if self._audio_type == 'csm': + from unsloth import FastModel + from transformers import CsmForConditionalGeneration + logger.info("Loading as CSM audio model...") + model, tokenizer = FastModel.from_pretrained( + model_name=checkpoint_path, + max_seq_length=max_seq_length, + dtype=None, + auto_model=CsmForConditionalGeneration, + load_in_4bit=False, + ) + + elif self._audio_type == 'whisper': + from unsloth import FastModel + from transformers import WhisperForConditionalGeneration + logger.info("Loading as Whisper audio model...") + model, tokenizer = FastModel.from_pretrained( + model_name=checkpoint_path, + dtype=None, + load_in_4bit=False, + auto_model=WhisperForConditionalGeneration, + ) + + elif self._audio_type == 'snac': + logger.info("Loading as SNAC (Orpheus) audio model...") + model, tokenizer = FastLanguageModel.from_pretrained( + model_name=checkpoint_path, + max_seq_length=max_seq_length, + dtype=None, + load_in_4bit=load_in_4bit, + ) + + elif self._audio_type == 'bicodec': + from unsloth import FastModel + logger.info("Loading as BiCodec (Spark-TTS) audio model...") + model, tokenizer = FastModel.from_pretrained( + model_name=checkpoint_path, + max_seq_length=max_seq_length, + dtype=torch.float32, + load_in_4bit=False, + ) + + elif self._audio_type == 'dac': + from unsloth import FastModel + logger.info("Loading as DAC (OuteTTS) audio model...") + model, tokenizer = FastModel.from_pretrained( + model_name=checkpoint_path, + max_seq_length=max_seq_length, + load_in_4bit=False, + ) + + elif self.is_vision: logger.info("Loading as vision model...") model, processor = FastVisionModel.from_pretrained( model_name=checkpoint_path, @@ -174,6 +228,7 @@ class ExportBackend: load_in_4bit=load_in_4bit, ) tokenizer = processor # For vision models, processor acts as tokenizer + else: logger.info("Loading as text model...") model, tokenizer = FastLanguageModel.from_pretrained( @@ -191,7 +246,12 @@ class ExportBackend: self.current_tokenizer = tokenizer self.current_checkpoint = checkpoint_path - model_type = "Vision" if self.is_vision else "Text" + if self._audio_type: + model_type = f"Audio ({self._audio_type})" + elif self.is_vision: + model_type = "Vision" + else: + model_type = "Text" peft_info = " (PEFT Adapter)" if self.is_peft else " (Merged Model)" logger.info(f"Successfully loaded {model_type} model{peft_info}") @@ -246,6 +306,9 @@ class ExportBackend: # Determine save method if format_type == "4-bit (FP4)": save_method = "merged_4bit_forced" + elif self._audio_type == 'whisper': + # Whisper uses save_method=None for local 16-bit merged save + save_method = None else: # 16-bit (FP16) save_method = "merged_16bit" @@ -271,10 +334,12 @@ class ExportBackend: logger.info(f"Pushing merged model to Hub: {repo_id}") + # Whisper uses save_method=None for local but "merged_16bit" for hub push + hub_save_method = save_method if save_method is not None else "merged_16bit" self.current_model.push_to_hub_merged( repo_id, self.current_tokenizer, - save_method=save_method, + save_method=hub_save_method, token=hf_token, private=private )