Spaces:
Running on Zero
Running on Zero
| """Text-only PitchIn-Q inference using the paper's merged checkpoint.""" | |
| import os | |
| import time | |
| from pathlib import Path | |
| import spaces # Must precede torch for ZeroGPU's CUDA initialization. | |
| import torch | |
| from peft import PeftModel | |
| from transformers import AutoConfig, AutoProcessor, Qwen2_5OmniForConditionalGeneration | |
| GENERATION = dict( | |
| do_sample=True, temperature=0.6, top_p=0.85, top_k=30, | |
| repetition_penalty=1.08, max_new_tokens=384, | |
| eos_token_id=151645, pad_token_id=151643, | |
| ) | |
| _ENGINE = None | |
| def generate_on_gpu(wave, question): | |
| # Keep the startup model in the forked worker's globals. Passing a bound | |
| # method's `self` through ZeroGPU's argument queue reserializes it on CPU. | |
| if _ENGINE is None: | |
| raise RuntimeError("PitchIn is not initialized.") | |
| return _ENGINE._generate(wave, question) | |
| class PitchIn: | |
| def __init__(self): | |
| global _ENGINE | |
| root = Path(__file__).parent | |
| merged_path = os.environ.get("PITCHIN_MODEL_PATH") | |
| model_path = merged_path or os.environ.get("PITCHIN_BASE_MODEL", "Qwen/Qwen2.5-Omni-7B") | |
| model_args = {"local_files_only": True} if Path(model_path).is_dir() else { | |
| "revision": "ae9e1690543ffd5c0221dc27f79834d0294cba00", | |
| } | |
| processor_path = merged_path or os.environ.get("PITCHIN_PROCESSOR_PATH", str(root / "processor")) | |
| config = AutoConfig.from_pretrained(model_path, **model_args) | |
| config.enable_audio_output = False | |
| self.processor = AutoProcessor.from_pretrained(processor_path, local_files_only=True) | |
| outer = Qwen2_5OmniForConditionalGeneration.from_pretrained( | |
| model_path, config=config, dtype=torch.bfloat16, | |
| attn_implementation="sdpa", **model_args, | |
| ) | |
| if not merged_path: | |
| adapter_path = os.environ.get("PITCHIN_ADAPTER_PATH", str(root / "adapter")) | |
| outer.thinker = PeftModel.from_pretrained( | |
| outer.thinker, adapter_path, is_trainable=False, torch_device="cpu", | |
| ).merge_and_unload() | |
| self.model = outer.thinker.to("cuda").eval() | |
| if self.processor.feature_extractor.sampling_rate != 16000: | |
| raise RuntimeError("Unexpected checkpoint audio sample rate.") | |
| _ENGINE = self | |
| def generate(self, wave, question): | |
| return generate_on_gpu(wave, question) | |
| def _generate(self, wave, question): | |
| started = time.monotonic() | |
| messages = [{"role": "user", "content": [ | |
| {"type": "audio", "audio": "uploaded_audio.wav"}, | |
| {"type": "text", "text": question}, | |
| ]}] | |
| text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) | |
| inputs = self.processor( | |
| text=[text], audio=[wave], sampling_rate=16000, return_tensors="pt", | |
| audio_kwargs={"truncation": False}, | |
| text_kwargs={"padding": True, "truncation": False}, | |
| ) | |
| inputs = {k: v.to("cuda", dtype=torch.bfloat16 if v.is_floating_point() else v.dtype) | |
| if isinstance(v, torch.Tensor) else v for k, v in inputs.items()} | |
| input_length = inputs["input_ids"].shape[1] | |
| with torch.inference_mode(): | |
| tokens = self.model.generate(**inputs, **GENERATION)[0, input_length:] | |
| response = self.processor.tokenizer.decode(tokens, skip_special_tokens=True).strip() | |
| if not response: | |
| raise RuntimeError("The model returned an empty answer. Please try again.") | |
| count = len(tokens) | |
| stopped_at_limit = count == GENERATION["max_new_tokens"] and tokens[-1].item() != GENERATION["eos_token_id"] | |
| status = f"{len(wave) / 16000:.1f}s audio · {time.monotonic() - started:.1f}s generation" | |
| if stopped_at_limit: | |
| status += " · Response reached the demo length limit." | |
| return response, status | |