Merge pull request #360 from Capsize-Games/develop

Offloads heavy functions to threads, prevents application from freezi…
Capsize-Games · Jan 10, 2024 · d4a8109 · d4a8109
2 parents 8f94753 + 5f2a1e9
commit d4a8109
Show file tree

Hide file tree

Showing 13 changed files with 587 additions and 569 deletions.
diff --git a/src/airunner/aihandler/runner.py b/src/airunner/aihandler/runner.py
@@ -970,17 +970,28 @@ def call_pipe(self, **kwargs):
             try:
                 args.update({
                     "prompt_embeds": self.prompt_embeds,
-                    "negative_prompt_embeds": self.negative_prompt_embeds,
                 })
             except Exception as _e:
                 Logger.warning("Compel failed: " + str(_e))
                 args.update({
                     "prompt": self.prompt,
-                    "negative_prompt": self.negative_prompt,
                 })
         else:
             args.update({
                 "prompt": self.prompt,
+            })
+        if self.use_compel:
+            try:
+                args.update({
+                    "negative_prompt_embeds": self.negative_prompt_embeds,
+                })
+            except Exception as _e:
+                Logger.warning("Compel failed: " + str(_e))
+                args.update({
+                    "negative_prompt": self.negative_prompt,
+                })
+        else:
+            args.update({
                 "negative_prompt": self.negative_prompt,
             })
         args["callback_steps"] = 1

diff --git a/src/airunner/aihandler/settings_manager.py b/src/airunner/aihandler/settings_manager.py
@@ -1,11 +1,11 @@
 from PyQt6.QtCore import QObject, pyqtSignal
 
 from airunner.aihandler.qtvar import StringVar, IntVar, BooleanVar, FloatVar, DictVar
-from airunner.data.db import session
-from airunner.data.models import LLMGenerator, Settings, GeneratorSetting, AIModel, Pipeline, ControlnetModel, ImageFilter, Prompt, \
-    SavedPrompt, PromptCategory, PromptVariable, PromptVariableCategory, PromptVariableCategoryWeight, StandardImageWidgetSettings
-from airunner.utils import save_session
+from airunner.data.models import LLMGenerator, Settings, GeneratorSetting, AIModel, Pipeline, ControlnetModel, ImageFilter, \
+    SavedPrompt, StandardImageWidgetSettings
+from airunner.utils import save_session, get_session
 from airunner.aihandler.logger import Logger as logger
+from airunner.data.models import Document
 
 document = None
 _app = None
@@ -17,6 +17,7 @@
     "FLOAT": FloatVar,
     "JSON": DictVar,
 }
+session = get_session()
 
 
 class SettingsSignal(QObject):
@@ -236,8 +237,7 @@ def __init__(self, app=None, *args, **kwargs):
             _app = app
             document = _app.document
         else:
-            from airunner.data.db import session
-            from airunner.data.models import Document
+            session = get_session()
             document = session.query(Document).first()
 
         super().__init__(*args, **kwargs)

diff --git a/src/airunner/aihandler/transformer_runner.py b/src/airunner/aihandler/transformer_runner.py
@@ -12,7 +12,7 @@
 
 from airunner.aihandler.settings_manager import SettingsManager
 from airunner.data.models import LLMGenerator
-from airunner.data.db import session
+from airunner.utils import get_session
 from airunner.aihandler.logger import Logger
 
 
@@ -66,6 +66,7 @@ class TransformerRunner(QObject):
     @property
     def generator(self):
         try:
+            session = get_session()
             if not self._generator or self.current_generator_name != self.requested_generator_name:
                 self.current_generator_name = self.requested_generator_name
                 self._generator = session.query(LLMGenerator).filter_by(name=self.current_generator_name).first()