joi/engine.py at main · hellopahe/joi · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
import threading
import queue
import numpy as np

from mlx_vlm import load
from mlx_vlm.models.qwen3_omni_moe.omni_utils import prepare_omni_inputs
from debug import log, is_debug


class InferenceEngine:
    def __init__(self, model_path, sample_rate=24000):
        self.model_path = model_path
        self.sample_rate = sample_rate
        self.model = None
        self.processor = None
        self.lock = threading.Lock()
        self.abort_flag = threading.Event()
        self.busy = False

    def load_model(self):
        log("ENGINE", "Loading model...")
        self.model, self.processor = load(self.model_path, trust_remote_code=True)
        log("ENGINE", "Model loaded")

    def is_busy(self):
        return self.busy

    def abort(self):
        self.abort_flag.set()

    def generate(self, conversation: list, output_queue: queue.Queue, **gen_config):
        self.abort_flag.clear()
        self.busy = True

        try:
            with self.lock:
                model_inputs, _ = prepare_omni_inputs(self.processor, conversation)

                if self.abort_flag.is_set():
                    output_queue.put(None)
                    return

                generate_kwargs = {
                    "input_ids": model_inputs["input_ids"],
                    "pixel_values": model_inputs.get("pixel_values"),
                    "pixel_values_videos": model_inputs.get("pixel_values_videos"),
                    "image_grid_thw": model_inputs.get("image_grid_thw"),
                    "video_grid_thw": model_inputs.get("video_grid_thw"),
                    "input_features": model_inputs.get("input_features"),
                    "feature_attention_mask": model_inputs.get("feature_attention_mask"),
                    "audio_feature_lengths": model_inputs.get("audio_feature_lengths"),
                    **gen_config,
                }
                generate_kwargs = {k: v for k, v in generate_kwargs.items() if v is not None}

                for event_type, data in self.model.generate_stream_interleaved(**generate_kwargs, debug=is_debug()):
                    if self.abort_flag.is_set():
                        break
                    if event_type == "text":
                        text = self.processor.decode(data)
                        output_queue.put(("text", text))
                    elif event_type == "audio":
                        audio_data = np.array(data.reshape(-1), dtype=np.float32)
                        output_queue.put(("audio", audio_data))

                output_queue.put(None)
        except Exception as e:
            log("ENGINE", f"Error: {e}")
            import traceback
            traceback.print_exc()
            output_queue.put(None)
        finally:
            self.busy = False