diff --git a/aiserver.py b/aiserver.py index 6315b3cd..a25967c1 100644 --- a/aiserver.py +++ b/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") diff --git a/environments/huggingface.yml b/environments/huggingface.yml index 25706993..18f8a8f3 100644 --- a/environments/huggingface.yml +++ b/environments/huggingface.yml @@ -22,6 +22,7 @@ dependencies: - loguru - Pillow - diffusers + - psutil - pip: - flask-cloudflared - flask-ngrok diff --git a/environments/rocm.yml b/environments/rocm.yml index c6bf06d2..d3c06fb3 100644 --- a/environments/rocm.yml +++ b/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.* diff --git a/gensettings.py b/gensettings.py index 04ee9e77..ae48979e 100644 --- a/gensettings.py +++ b/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 =[{ diff --git a/koboldai_settings.py b/koboldai_settings.py index f066bcbf..4fb5a937 100644 --- a/koboldai_settings.py +++ b/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 diff --git a/requirements.txt b/requirements.txt index 02972a2e..8aee7dcf 100644 --- a/requirements.txt +++ b/requirements.txt @@ -17,4 +17,5 @@ marshmallow>=3.13 apispec-webframeworks loguru Pillow -diffusers \ No newline at end of file +diffusers +psutil \ No newline at end of file diff --git a/requirements_mtj.txt b/requirements_mtj.txt index 12af394d..3df1b04a 100644 --- a/requirements_mtj.txt +++ b/requirements_mtj.txt @@ -21,4 +21,5 @@ marshmallow>=3.13 apispec-webframeworks loguru Pillow -diffusers \ No newline at end of file +diffusers +psutil \ No newline at end of file diff --git a/templates/settings flyout.html b/templates/settings flyout.html index 49895852..9c61676b 100644 --- a/templates/settings flyout.html +++ b/templates/settings flyout.html @@ -276,6 +276,16 @@ +
+

expand_more Images

+
+
+ {% with menu='Interface' %} + {% with sub_path='Images' %} + {% include 'settings item.html' %} + {% endwith %} + {% endwith %} +

expand_more Formatting