Better image generation

aiserver.py

@@ -8180,7 +8180,7 @@ def UI_2_generate_image(data):
     #If we have > 4 keys, use those otherwise use sumarization
     if len(keys) < 4:
         from transformers import pipeline as summary_pipeline
-        summarizer = summary_pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
+        summarizer = summary_pipeline("summarization", model="sshleifer/distilbart-cnn-12-1")
         #text to summarize:
         if len(koboldai_vars.actions) < 5:
             text = "".join(koboldai_vars.actions[:-5]+[koboldai_vars.prompt])
@@ -8191,24 +8191,27 @@ def UI_2_generate_image(data):
         transformers.generation_utils.GenerationMixin._get_stopping_criteria = old_transfomers_functions['transformers.generation_utils.GenerationMixin._get_stopping_criteria']
         keys = [summarizer(text, max_length=100, min_length=30, do_sample=False)[0]['summary_text']]
         transformers.generation_utils.GenerationMixin._get_stopping_criteria = temp
-        del summarizer
     
     art_guide = 'fantasy illustration, artstation, by jason felix by steve argyle by tyler jacobson by peter mohrbacher, cinematic lighting', 
 
-    #If we don't have a GPU, use horde
-    if not koboldai_vars.hascuda:
-        b64_data = text2img(", ".join(keys), art_guide = art_guide)
+    #If we don't have a GPU, use horde if we're allowed to
+    if (not koboldai_vars.hascuda and koboldai_vars.img_gen_priority != 0) or  koboldai_vars.img_gen_priority == 3:
+        b64_data = text2img_horde(", ".join(keys), art_guide = art_guide)
     else:
+        import psutil
+        #We aren't being forced to use horde, so now let's figure out if we should use local
         if torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_reserved(0) >= 6000000000:
             #He have enough vram, just do it locally
             b64_data = text2img_local(", ".join(keys), art_guide = art_guide)
-        elif torch.cuda.get_device_properties(0).total_memory > 6000000000:
+        elif torch.cuda.get_device_properties(0).total_memory > 6000000000 and koboldai_vars.img_gen_priority <= 1:
             #We could do it locally by swapping the model out
             print("Could do local or online")
-        else:
-            b64_data = text2img_local(", ".join(keys), art_guide = art_guide)
+            b64_data = text2img_horde(", ".join(keys), art_guide = art_guide)
+        elif koboldai_vars.img_gen_priority != 0:
+            b64_data = text2img_horde(", ".join(keys), art_guide = art_guide)
     koboldai_vars.picture = b64_data
     koboldai_vars.picture_prompt = ", ".join(keys)
+    koboldai_vars.generating_image = False
     #emit("Action_Image", {'b64': b64_data, 'prompt': ", ".join(keys)})
     
 
@@ -8216,11 +8219,15 @@ def UI_2_generate_image(data):
 def text2img_local(prompt, art_guide="", filename="new.png"):
     start_time = time.time()
     print("Generating Image")
+    koboldai_vars.aibusy = True
     koboldai_vars.generating_image = True
     from diffusers import StableDiffusionPipeline
     import base64
     from io import BytesIO
-    pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16, cache="./stable-diffusion-v1-4").to("cuda")
+    if koboldai_vars.image_pipeline is None:
+        pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16, cache="./stable-diffusion-v1-4").to("cuda")
+    else:
+        pipe = koboldai_vars.image_pipeline
     print("time to load: {}".format(time.time() - start_time))
     from torch import autocast
     with autocast("cuda"):
@@ -8229,13 +8236,20 @@ def text2img_local(prompt, art_guide="", filename="new.png"):
     image.save(buffered, format="JPEG")
     img_str = base64.b64encode(buffered.getvalue()).decode('ascii')
     print("time to generate: {}".format(time.time() - start_time))
-    pipe.to("cpu")
+    if koboldai_vars.keep_img_gen_in_memory:
+        pipe.to("cpu")
+        if koboldai_vars.image_pipeline is None:
+        koboldai_vars.image_pipeline = pipe
+    else:
+        koboldai_vars.image_pipeline = None
+        del pipe
     koboldai_vars.generating_image = False
+    koboldai_vars.aibusy = True
     print("time to unload: {}".format(time.time() - start_time))
     return img_str
 
 @logger.catch
-def text2img(prompt, 
+def text2img_horde(prompt, 
              art_guide = 'fantasy illustration, artstation, by jason felix by steve argyle by tyler jacobson by peter mohrbacher, cinematic lighting', 
              filename = "story_art.png"):
     print("Generating Image using Horde")

environments/huggingface.yml

@@ -22,6 +22,7 @@ dependencies:
   - loguru
   - Pillow
   - diffusers
+  - psutil
   - pip:
     - flask-cloudflared
     - flask-ngrok

environments/rocm.yml

@@ -19,6 +19,7 @@ dependencies:
   - loguru
   - Pillow
   - diffusers
+  - psutil
   - pip:
     - --find-links https://download.pytorch.org/whl/rocm4.2/torch_stable.html
     - torch==1.10.*

gensettings.py

@@ -522,6 +522,36 @@ gensettingstf = [
     "classname": "user",
     "name": "beep_on_complete"
  	},
+    {
+    "UI_V2_Only": True,
+ 	"uitype": "dropdown",
+	"unit": "text",
+	"label": "Image Priority",
+	"id": "img_gen_priority",
+	"default": 1,
+    "tooltip": "Determine where images will be generated.",
+    "menu_path": "Interface",
+    "sub_path": "Images",
+    "classname": "user",
+    "name": "img_gen_priority",
+    'children': [{'text': 'Use Local Only', 'value': 0}, {'text':'Perfer Local','value':1}, {'text':'Perfer Horde', 'value':2}, {'text':'Use Horde Only', 'value':3}]
+ 	},
+    {
+    "UI_V2_Only": True,
+ 	"uitype": "toggle",
+ 	"unit": "bool",
+ 	"label": "Model in memory",
+ 	"id": "keep_img_gen_in_memory",
+ 	"min": 0,
+ 	"max": 1,
+ 	"step": 1,
+ 	"default": 0,
+	"tooltip": "If enabled, the system will keep the model in memory speeding up image generation times",
+    "menu_path": "Interface",
+    "sub_path": "Images",
+    "classname": "user",
+    "name": "keep_img_gen_in_memory"
+ 	},
 ]
 
 gensettingsik =[{

koboldai_settings.py

@@ -714,6 +714,8 @@ class user_settings(settings):
         self.output_streaming = True
         self.show_probs = False # Whether or not to show token probabilities
         self.beep_on_complete = False
+        self.img_gen_priority = 1
+        self.keep_img_gen_in_memory = False
         
         
     def __setattr__(self, name, value):
@@ -725,8 +727,8 @@ class user_settings(settings):
             process_variable_changes(self.socketio, self.__class__.__name__.replace("_settings", ""), name, value, old_value)
         
 class system_settings(settings):
-    local_only_variables = ['socketio', 'lua_state', 'lua_logname', 'lua_koboldbridge', 'lua_kobold', 'lua_koboldcore', 'regex_sl', 'acregex_ai', 'acregex_ui', 'comregex_ai', 'comregex_ui', 'sp', '_horde_pid']
-    no_save_variables = ['socketio', 'lua_state', 'lua_logname', 'lua_koboldbridge', 'lua_kobold', 'lua_koboldcore', 'sp', '_horde_pid', 'horde_share', 'aibusy', 'serverstarted']
+    local_only_variables = ['socketio', 'lua_state', 'lua_logname', 'lua_koboldbridge', 'lua_kobold', 'lua_koboldcore', 'regex_sl', 'acregex_ai', 'acregex_ui', 'comregex_ai', 'comregex_ui', 'sp', '_horde_pid', 'image_pipeline']
+    no_save_variables = ['socketio', 'lua_state', 'lua_logname', 'lua_koboldbridge', 'lua_kobold', 'lua_koboldcore', 'sp', '_horde_pid', 'horde_share', 'aibusy', 'serverstarted', 'image_pipeline']
     settings_name = "system"
     def __init__(self, socketio):
         self.socketio = socketio
@@ -804,6 +806,7 @@ class system_settings(settings):
         self._horde_pid = None
         self.sh_apikey   = ""     # API key to use for txt2img from the Stable Horde.
         self.generating_image = False #The current status of image generation
+        self.image_pipeline = None
         self.cookies = {} #cookies for colab since colab's URL changes, cookies are lost
         
         

requirements.txt

@@ -17,4 +17,5 @@ marshmallow>=3.13
 apispec-webframeworks
 loguru
 Pillow
-diffusers
+diffusers
+psutil

requirements_mtj.txt

@@ -21,4 +21,5 @@ marshmallow>=3.13
 apispec-webframeworks
 loguru
 Pillow
-diffusers
+diffusers
+psutil

templates/settings flyout.html

@@ -276,6 +276,16 @@
 				<span class="setting_maxlabel"><span style="top: -4px; position: relative;"></span></span>
 			</div>
 		</div>
+		<div class="collapsable_header" onclick="toggle_setting_category(this);">
+			<h4 style="width:var(--flyout_menu_width);"><span class="material-icons-outlined cursor">expand_more</span> Images</h4>
+		</div>
+		<div class="setting_tile_area">
+			{% with menu='Interface' %}
+				{% with sub_path='Images' %}
+					{% include 'settings item.html' %}
+				{% endwith %}
+			{% endwith %}
+		</div>
 		<div class="collapsable_header" onclick="toggle_setting_category(this);">
 			<h4 style="width:var(--flyout_menu_width);"><span class="material-icons-outlined cursor">expand_more</span> Formatting</h4>
 		</div>