1from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor
2from qwen_omni_utils import process_mm_info
3from peft import PeftModel
4import torch
5def merge_lora(
6 base_model_path: str,
7 lora_checkpoint_path: str,
8 extra_file: str = "spk_dict.pt",
9 submodule_name: str = "thinker",
10 cache_dir: str = "cache/"
11):
12 """Load the original model, merge the LoRA weights.
13 For a specified submodule, and save the final merged model along with its configurations.
14 Args:
15 base_model_path (str): Path to the original model directory.
16 lora_checkpoint_path (str): Path to the directory containing LoRA weights.
17 extra_file (str): Name of the extra file to be copied (default: "spk_dict.pt").
18 submodule_name (str): Name of the submodule to merge (default: "thinker").
19 save_path (str): Directory where the merged model and configurations will be saved.
20 """
21 # 1. Load the original model
22 model = Qwen2_5OmniForConditionalGeneration.from_pretrained(base_model_path, torch_dtype=torch.bfloat16, device_map="auto", cache_dir=cache_dir)
23 print("Successfully loaded the original model.")
24 # 2. Extract the submodule to be merged (e.g., model.thinker)
25 if not hasattr(model, submodule_name):
26 raise AttributeError(f"The model does not have a submodule named '{submodule_name}'.")
27 base_submodule = getattr(model, submodule_name)
28 print(f"Successfully extracted submodule: {submodule_name}.")
29 # 3. Load the LoRA weights onto the extracted submodule
30 lora_model = PeftModel.from_pretrained(base_submodule, lora_checkpoint_path)
31 processor = Qwen2_5OmniProcessor.from_pretrained(lora_checkpoint_path)
32 print("LoRA weights and processor loaded successfully.")
33 # 4. Merge the LoRA weights into the submodule and unload the LoRA modules
34 merged_submodule = lora_model.merge_and_unload()
35 print("LoRA weights merged successfully.")
36 # 5. Replace the original submodule with the merged submodule in the model
37 setattr(model, submodule_name, merged_submodule)
38 return model
39cache_dir = "" # SET YOUR HF CACHE DIR
40assert cache_dir != ""
41model = merge_lora("Qwen/Qwen2.5-Omni-7B", "kit-isl-ai4lt/qwen_omni_lt_v1", cache_dir)
42model.disable_talker() # By default it is enabled
43### FOR BASE MODEL
44### model = Qwen2_5OmniForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-Omni-7B", torch_dtype=torch.bfloat16, device_map="auto", cache_dir=cache_dir)
45print("Disabled Talker for efficiency and is not working after LoRA text output fine-tuning")
46processor = Qwen2_5OmniProcessor.from_pretrained("Qwen/Qwen2.5-Omni-7B", cache_dir = cache_dir)
47# You can directly insert a local file path, a URL, or a base64-encoded image into the position where you want in the text.
48conversation = [
49 {
50 "role": "system",
51 "content": [
52 {"type": "text", "text": "You are a helpful assistant that translates audio into text."}
53 ],
54 },
55 {
56 "role": "user",
57 "content": [
58 {"type": "audio", "audio": "./allo.mp3"},
59 {"type": "text", "text": "Translate into English:"}
60 ],
61 },
62]
63text = processor.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
64audios, images, videos = process_mm_info(conversation, use_audio_in_video=True)
65inputs = processor(text=text, images=images, videos=videos, audio=audios, padding=True, return_tensors="pt").to(model.device)
66generated_ids = model.generate(**inputs)
67### Print only generated text with slicing
68print(processor.batch_decode(generated_ids[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)[0])
