Merge pull request #395 from Capsize-Games/develop

Code maintenence

src/airunner/aihandler/base_handler.py

@@ -0,0 +1,13 @@
+from PyQt6.QtCore import QObject
+
+from airunner.mediator_mixin import MediatorMixin
+from airunner.windows.main.settings_mixin import SettingsMixin
+from airunner.aihandler.logger import Logger
+
+
+class BaseHandler(QObject, MediatorMixin, SettingsMixin):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        SettingsMixin.__init__(self)
+        MediatorMixin.__init__(self)
+        self.logger = Logger(prefix=self.__class__.__name__)

src/airunner/aihandler/llm.py → src/airunner/aihandler/llm_handler.py

@@ -9,15 +9,14 @@
 from transformers import InstructBlipForConditionalGeneration
 from transformers import InstructBlipProcessor
 from transformers import TextIteratorStreamer
-
-from PyQt6.QtCore import QObject
+from airunner.aihandler.base_handler import BaseHandler
 
 from airunner.aihandler.logger import Logger
 from airunner.mediator_mixin import MediatorMixin
 
 
-class LLM(QObject, MediatorMixin):
-    logger = Logger(prefix="LLM")
+class LLMHandler(BaseHandler):
+    logger = Logger(prefix="LLMHandler")
     dtype = ""
     local_files_only = True
     set_attention_mask = False
@@ -78,10 +77,6 @@ def has_gpu(self):
         if self.dtype == "32bit" or not self.use_gpu:
             return False
         return torch.cuda.is_available()
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        MediatorMixin.__init__(self)
 
     def move_to_cpu(self):
         if self.model:

src/airunner/aihandler/runner.py → src/airunner/aihandler/sd_handler.py

@@ -22,6 +22,7 @@
 from diffusers import StableDiffusionControlNetPipeline, StableDiffusionControlNetImg2ImgPipeline, StableDiffusionControlNetInpaintPipeline, AsymmetricAutoencoderKL
 from diffusers import ConsistencyDecoderVAE
 from transformers import AutoFeatureExtractor
+from airunner.aihandler.base_handler import BaseHandler
 
 from airunner.aihandler.enums import FilterType
 from airunner.aihandler.mixins.compel_mixin import CompelMixin
@@ -40,16 +41,14 @@
 from airunner.windows.main.pipeline_mixin import PipelineMixin
 from airunner.windows.main.controlnet_model_mixin import ControlnetModelMixin
 from airunner.windows.main.ai_model_mixin import AIModelMixin
-from airunner.mediator_mixin import MediatorMixin
-from airunner.windows.main.settings_mixin import SettingsMixin
 from airunner.service_locator import ServiceLocator
 from airunner.utils import clear_memory
 
 torch.backends.cuda.matmul.allow_tf32 = True
 
 
-class SDRunner(
-    QObject,
+class SDHandler(
+    BaseHandler,
     MergeMixin,
     LoraMixin,
     MemoryEfficientMixin,
@@ -59,16 +58,14 @@ class SDRunner(
     SchedulerMixin,
 
     # Data Mixins
-    SettingsMixin,
     LayerMixin,
     LoraDataMixin,
     EmbeddingDataMixin,
     PipelineMixin,
     ControlnetModelMixin,
     AIModelMixin,
-    MediatorMixin
 ):
-    logger = Logger(prefix="SDRunner")
+    logger = Logger(prefix="SDHandler")
     _current_model: str = ""
     _previous_model: str = ""
     _initialized: bool = False
@@ -720,15 +717,13 @@ def original_model_data(self):
 
     def  __init__(self, **kwargs):
         #self.logger.set_level(LOG_LEVEL)
-        MediatorMixin.__init__(self)
-        SettingsMixin.__init__(self)
+        super().__init__()
         LayerMixin.__init__(self)
         LoraDataMixin.__init__(self)
         EmbeddingDataMixin.__init__(self)
         PipelineMixin.__init__(self)
         ControlnetModelMixin.__init__(self)
         AIModelMixin.__init__(self)
-        super().__init__()
         self.logger.info("Loading Stable Diffusion model runner...")
         self.safety_checker_model = self.models_by_pipeline_action("safety_checker")
         self.text_encoder_model = self.models_by_pipeline_action("text_encoder")

src/airunner/aihandler/speech_to_text.py → src/airunner/aihandler/stt_handler.py

@@ -3,15 +3,12 @@
 
 from transformers import AutoProcessor, WhisperForConditionalGeneration, AutoFeatureExtractor
 
-from PyQt6.QtCore import QObject
+from airunner.aihandler.base_handler import BaseHandler
+
 from PyQt6.QtCore import pyqtSignal
 
-from airunner.aihandler.logger import Logger
-from airunner.mediator_mixin import MediatorMixin
-
 
-class SpeechToText(QObject, MediatorMixin):
-    logger = Logger(prefix="SpeechToText")
+class STTHandler(BaseHandler):
     listening = False
     move_to_cpu_signal = pyqtSignal()
 
@@ -26,9 +23,8 @@ def on_move_to_cpu(self):
         self.logger.info("Moving model to CPU")
         self.model = self.model.to("cpu")
 
-    def __init__(self):
-        super().__init__()
-        MediatorMixin.__init__(self)
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
         self.load_model()
         self.register("move_to_cpu_signal", self)
         self.register("process_audio", self)

src/airunner/aihandler/tts.py → src/airunner/aihandler/tts_handler.py

@@ -3,25 +3,21 @@
 
 from queue import Queue
 
-from PyQt6.QtCore import QObject, pyqtSlot
+from PyQt6.QtCore import pyqtSlot
 
 from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan, BarkModel, BarkProcessor
 from datasets import load_dataset
+from airunner.aihandler.base_handler import BaseHandler
 
-from airunner.aihandler.logger import Logger
-from airunner.mediator_mixin import MediatorMixin
-from airunner.windows.main.settings_mixin import SettingsMixin
 
-
-class TTS(QObject, MediatorMixin, SettingsMixin):
+class TTSHandler(BaseHandler):
     """
     Generates speech from given text. 
     Responsible for managing the model, processor, vocoder, and speaker embeddings.
     Generates using either the SpeechT5 or Bark model.
 
     Use from a worker to avoid blocking the main thread.
     """
-    logger = Logger(prefix="TTS")
     character_replacement_map = {
         "\n": " ",
         "’": "'",
@@ -130,9 +126,7 @@ def sentence_chunks(self):
         return self.tts_settings["sentence_chunks"]
 
     def __init__(self, *args, **kwargs):
-        super().__init__()
-        SettingsMixin.__init__(self)
-        MediatorMixin.__init__(self)
+        super().__init__(*args, **kwargs)
         self.logger.info("Loading")
         self.corpus = []
         self.processor = None

src/airunner/aihandler/vision_handler.py

@@ -0,0 +1,8 @@
+from PyQt6.QtCore import QObject
+from airunner.aihandler.base_handler import BaseHandler
+
+from airunner.mediator_mixin import MediatorMixin
+
+
+class VisionHandler(BaseHandler):
+    pass

src/airunner/windows/main/main_window.py

@@ -37,7 +37,7 @@
 from airunner.windows.main.controlnet_model_mixin import ControlnetModelMixin
 from airunner.windows.main.ai_model_mixin import AIModelMixin
 from airunner.windows.main.image_filter_mixin import ImageFilterMixin
-from airunner.aihandler.engine import Engine
+from airunner.worker_manager import WorkerManager
 from airunner.mediator_mixin import MediatorMixin
 from airunner.service_locator import ServiceLocator
 
@@ -300,7 +300,7 @@ def __init__(self, *args, **kwargs):
         ServiceLocator.register("get_callback_for_slider", self.get_callback_for_slider)
 
 
-        self.engine = Engine()
+        self.engine = WorkerManager()
 
         self.ui.setupUi(self)
 

src/airunner/aihandler/engine.py → src/airunner/worker_manager.py

@@ -16,10 +16,11 @@
 from airunner.workers.sd_generate_worker import SDGenerateWorker
 from airunner.workers.sd_request_worker import SDRequestWorker
 from airunner.aihandler.logger import Logger
-from airunner.aihandler.tts import TTS
 from airunner.windows.main.settings_mixin import SettingsMixin
 from airunner.service_locator import ServiceLocator
 from airunner.utils import clear_memory
+from airunner.workers.vision_capture_worker import VisionCaptureWorker
+from airunner.workers.vision_processor_worker import VisionProcessorWorker
 
 
 class Message:
@@ -29,7 +30,7 @@ def __init__(self, *args, **kwargs):
         self.conversation = kwargs.get("conversation")
 
 
-class Engine(QObject, MediatorMixin, SettingsMixin):
+class WorkerManager(QObject, MediatorMixin, SettingsMixin):
     """
     The engine is responsible for processing requests and offloading
     them to the appropriate AI model controller.
@@ -51,12 +52,12 @@ def do_response(self, response):
         Handle a response from the application by putting it into
         a response worker queue.
         """
-        self.response_worker.add_to_queue(response)
+        self.engine_response_worker.add_to_queue(response)
 
     def on_engine_cancel_signal(self, _ignore):
         self.logger.info("Canceling")
         self.emit("sd_cancel_signal")
-        self.request_worker.cancel()
+        self.engine_request_worker.cancel()
 
     def on_engine_stop_processing_queue_signal(self):
         self.do_process_queue = False
@@ -87,11 +88,6 @@ def __init__(self, **kwargs):
         self.logger = Logger(prefix="Engine")
         self.clear_memory()
 
-        # Initialize Controllers
-        #self.stt_controller = STTController(engine=self)
-        # self.ocr_controller = ImageProcessor(engine=self)
-        self.tts_controller = TTS(engine=self)
-
         self.register("hear_signal", self)
         self.register("engine_cancel_signal", self)
         self.register("engine_stop_processing_queue_signal", self)
@@ -107,30 +103,32 @@ def __init__(self, **kwargs):
         self.register("image_generate_request_signal", self)
         self.register("llm_response_signal", self)
         self.register("llm_text_streamed_signal", self)
+        self.register("AudioCaptureWorker_response_signal", self)
+        self.register("AudioProcessorWorker_processed_audio", self)
 
         self.sd_request_worker = self.create_worker(SDRequestWorker)
         self.sd_generate_worker = self.create_worker(SDGenerateWorker)
 
-        self.request_worker = self.create_worker(EngineRequestWorker)
-        self.response_worker = self.create_worker(EngineResponseWorker)
+        self.engine_request_worker = self.create_worker(EngineRequestWorker)
+        self.engine_response_worker = self.create_worker(EngineResponseWorker)
 
-        self.generator_worker = self.create_worker(TTSGeneratorWorker)
-        self.vocalizer_worker = self.create_worker(TTSVocalizerWorker)
+        self.tts_generator_worker = self.create_worker(TTSGeneratorWorker)
+        self.tts_vocalizer_worker = self.create_worker(TTSVocalizerWorker)
 
-        self.request_worker = self.create_worker(LLMRequestWorker)
-        self.generate_worker = self.create_worker(LLMGenerateWorker)
+        self.llm_request_worker = self.create_worker(LLMRequestWorker)
+        self.llm_generate_worker = self.create_worker(LLMGenerateWorker)
 
-        self.audio_capture_worker = self.create_worker(AudioCaptureWorker)
-        self.audio_processor_worker = self.create_worker(AudioProcessorWorker)
+        self.stt_audio_capture_worker = self.create_worker(AudioCaptureWorker)
+        self.stt_audio_processor_worker = self.create_worker(AudioProcessorWorker)
 
-        self.register("AudioCaptureWorker_response_signal", self)
-        self.register("AudioProcessorWorker_processed_audio", self)
+        self.vision_capture_worker = self.create_worker(VisionCaptureWorker)
+        self.vision_processor_worker = self.create_worker(VisionProcessorWorker)
 
         self.register("tts_request", self)
 
     def on_AudioCaptureWorker_response_signal(self, message: np.ndarray):
         self.logger.info("Heard signal")
-        self.audio_processor_worker.add_to_queue(message)
+        self.stt_audio_processor_worker.add_to_queue(message)
 
     def on_AudioProcessorWorker_processed_audio(self, message: np.ndarray):
         self.logger.info("Processed audio")
@@ -140,7 +138,7 @@ def on_LLMGenerateWorker_response_signal(self, message:dict):
         self.emit("llm_response_signal", message)
 
     def on_tts_request(self, data: dict):
-        self.generator_worker.add_to_queue(data)
+        self.tts_generator_worker.add_to_queue(data)
 
     def on_llm_response_signal(self, message):
         self.do_response(message)
@@ -195,7 +193,7 @@ def do_image_generate_request(self, message):
         ))
 
     def request_queue_size(self):
-        return self.request_worker.queue.qsize()
+        return self.engine_request_worker.queue.qsize()
 
     def do_listen(self):
         # self.stt_controller.do_listen()
@@ -227,8 +225,8 @@ def on_clear_llm_history_signal(self):
 
     def stop(self):
         self.logger.info("Stopping")
-        self.request_worker.stop()
-        self.response_worker.stop()
+        self.engine_request_worker.stop()
+        self.engine_response_worker.stop()
 
     def move_sd_to_cpu(self):
         if ServiceLocator.get("is_pipe_on_cpu")() or not ServiceLocator.get("has_pipe")():

src/airunner/workers/audio_processor_worker.py

@@ -1,4 +1,4 @@
-from airunner.aihandler.speech_to_text import SpeechToText
+from airunner.aihandler.stt_handler import STTHandler
 from airunner.workers.worker import Worker
 
 
@@ -11,7 +11,7 @@ class AudioProcessorWorker(Worker):
 
     def __init__(self, prefix):
         super().__init__(prefix=prefix)
-        self.stt = SpeechToText()
+        self.stt = STTHandler()
         self.register("stt_audio_processed", self)
 
     def on_stt_audio_processed(self, transcription):

src/airunner/workers/llm_generate_worker.py

@@ -1,10 +1,10 @@
 from airunner.workers.worker import Worker
-from airunner.aihandler.llm import LLM
+from airunner.aihandler.llm_handler import LLMHandler
 
 
 class LLMGenerateWorker(Worker):
     def __init__(self, prefix="LLMGenerateWorker"):
-        self.llm = LLM()
+        self.llm = LLMHandler()
         super().__init__(prefix=prefix)
         self.register("clear_history", self)
         self.register("LLMRequestWorker_response_signal", self)

src/airunner/workers/sd_generate_worker.py

@@ -3,15 +3,15 @@
 
 from airunner.aihandler.enums import EngineResponseCode
 from airunner.workers.worker import Worker
-from airunner.aihandler.runner import SDRunner
+from airunner.aihandler.sd_handler import SDHandler
 
 torch.backends.cuda.matmul.allow_tf32 = True
 
 
 class SDGenerateWorker(Worker):
     def __init__(self, prefix="SDGenerateWorker"):
         super().__init__(prefix=prefix)
-        self.sd = SDRunner()
+        self.sd = SDHandler()
         self.register("add_sd_response_to_queue_signal", self)
 
     def on_add_sd_response_to_queue_signal(self, request):

src/airunner/workers/tts_generator_worker.py

@@ -1,7 +1,7 @@
 import time
 
 from airunner.workers.worker import Worker
-from airunner.aihandler.tts import TTS
+from airunner.aihandler.tts_handler import TTSHandler
 
 
 class TTSGeneratorWorker(Worker):
@@ -10,7 +10,7 @@ class TTSGeneratorWorker(Worker):
     """
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.tts = TTS()
+        self.tts = TTSHandler()
         self.tts.run()
         self.play_queue = []
         self.play_queue_started = False
@@ -49,7 +49,6 @@ def generate(self, message):
         else:
             response = self.generate_with_t5(text)
 
-        print("adding to stream", response)
         self.emit("TTSGeneratorWorker_add_to_stream_signal", response)
 
     def move_inputs_to_device(self, inputs):