forked from glibsonoran/Plush-for-ComfyUI
-
Notifications
You must be signed in to change notification settings - Fork 0
/
style_prompt.py
1359 lines (1071 loc) · 55.3 KB
/
style_prompt.py
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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
import os
import base64
from io import BytesIO
from PIL import Image, ImageOps, TiffImagePlugin, UnidentifiedImageError
import folder_paths
import numpy as np
import torch
from typing import Optional
from enum import Enum
import requests
from requests.adapters import HTTPAdapter, Retry
from openai import OpenAI
import anthropic
from .mng_json import json_manager, helpSgltn, TroubleSgltn # add .
from . import api_requests as rqst
from .fetch_models import FetchModels, ModelUtils, RequestMode # add .
#pip install pillow
#pip install bytesio
#Enum for style_prompt user input modes
class InputMode(Enum):
IMAGE_PROMPT = 1
IMAGE_ONLY = 2
PROMPT_ONLY = 3
#Get information from the config.json file
class cFigSingleton:
_instance = None
def __new__(cls):
if cls._instance is None:
cls._instance = super().__new__(cls)
cls._lm_client = None
cls._anthropic_client = None
cls._lm_url = ""
cls._lm_request_mode = None
cls._lm_key = ""
cls._groq_key = ""
cls._claude_key = ""
cls._gemini_key = ""
cls._lm_models = None
cls._groq_models = None
cls._claude_models = None
cls._gemini_models = None
cls._ollama_models = None
cls._optional_models = None
cls._written_url = ""
cls.j_mngr = json_manager()
cls._model_fetch = FetchModels()
cls._model_prep = ModelUtils()
cls._pyexiv2 = None
cls._instance.get_file()
return cls._instance
def get_file(self):
#Get script working directory
#j_mngr = json_manager()
# Error handling is in the load_json method
# Errors will be raised since is_critical is set to True
config_data =self.j_mngr.load_json(self.j_mngr.config_file, True)
#Pyexiv2 seems to have trouble loading with some Python versions (it's misreading the vesrion number)
#So I'll open it in a try block so as not to stop the whole suite from loading
try:
import pyexiv2
self._pyexiv2 = pyexiv2
except Exception as e:
self._pyexiv2 = None
self.j_mngr.log_events(f"The Pyexiv2 library failed to load with Error: {e} ",
TroubleSgltn.Severity.ERROR)
#check if file is empty
if not config_data:
raise ValueError("Plush - Error: config.json contains no valid JSON data")
# Try getting API key from Plush environment variable
self._fig_key = os.getenv('OAI_KEY',"") or os.getenv('OPENAI_API_KEY',"")
if not self._fig_key:
#Let user know some nodes will not function
self.j_mngr.log_events("Open AI API key invalid or not found, some nodes will not be functional. See Read Me to install the key",
TroubleSgltn.Severity.WARNING)
self._lm_key = os.getenv("LLM_KEY","") #Fetch the LLM_KEY if the user has created one
self._groq_key = os.getenv("GROQ_API_KEY", "") or os.getenv("LLM_KEY", "")
self._claude_key = os.getenv("ANTHROPIC_API_KEY", "") or os.getenv("LLM_KEY", "")
self._gemini_key = os.getenv("GEMINI_API_KEY", "") or os.getenv("LLM_KEY", "")
#Get user saved Open Source URL from the text file
#At this point all this does is pre-populate new instances of the node.
url_file =self.j_mngr.append_filename_to_path(self.j_mngr.script_dir, 'OpenSourceURL.txt')
lm_url_list =self.j_mngr.read_lines_of_file(url_file)
if lm_url_list:
self._lm_url = lm_url_list[0] #This call will set up the client if the local LLM server is running
self._written_url = self._lm_url
self.figInstruction = config_data.get('instruction', "")
self.figExample = config_data.get('example', "")
self.figExample2 = config_data.get('example2', "")
self.fig_n_example = config_data.get('n_example', "")
self.fig_n_example2 = config_data.get('n_example2', "")
self._use_examples = False
self.figStyle = config_data.get('style', "")
self.figImgInstruction = config_data.get('img_instruction', "")
self.figImgPromptInstruction = config_data.get('img_prompt_instruction', "")
self.fig_n_Instruction = config_data.get('n_instruction', "")
self.fig_n_ImgPromptInstruction = config_data.get('n_img_prompt_instruction', "")
self.fig_n_ImgInstruction = config_data.get('n_img_instruction', "")
self._fig_gpt_models = []
if self._fig_key:
try:
self.figOAIClient = OpenAI(api_key= self._fig_key)
except Exception as e:
self.j_mngr.log_events(f"Invalid or missing OpenAI API key. Please note, keys must now be kept in an environment variable (see: ReadMe) {e}",
severity=TroubleSgltn.Severity.ERROR)
if self._claude_key:
try:
self._anthropic_client = anthropic.Anthropic(api_key = self._claude_key)
except Exception as e:
self.j_mngr.log_events(f"Invalid or missing Anthropic API key. Please note, keys must be kept in an environment variable.{e}",
severity=TroubleSgltn.Severity.ERROR)
self._fig_gpt_models = self._model_fetch.fetch_models(RequestMode.OPENAI, self._fig_key)
self._groq_models = self._model_fetch.fetch_models(RequestMode.GROQ, self._groq_key)
self._claude_models = self._model_fetch.fetch_models(RequestMode.CLAUDE, self._claude_key)
self._gemini_models = self._model_fetch.fetch_models(RequestMode.GEMINI, self._gemini_key)
self._ollama_models = self._model_fetch.fetch_models(RequestMode.OLLAMA, "")
self._optional_models = self._model_fetch.fetch_models(RequestMode.OPENSOURCE, "")
def get_chat_models(self, sort_it:bool=False, filter_str:tuple=())->list:
return self._model_prep.prep_models_list(self._fig_gpt_models, sort_it, filter_str)
def get_groq_models(self, sort_it:bool=False, filter_str:tuple=()):
return self._model_prep.prep_models_list(self._groq_models, sort_it, filter_str)
def get_claude_models(self, sort_it:bool=False, filter_str:tuple=())->list:
return self._model_prep.prep_models_list(self._claude_models, sort_it, filter_str)
def get_gemini_models(self, sort_it:bool=False, filter_str:tuple=())->list:
return self._model_prep.prep_models_list(self._gemini_models, sort_it, filter_str)
def get_ollama_models(self, sort_it:bool=False, filter_str:tuple=())->list:
return self._model_prep.prep_models_list(self._ollama_models, sort_it, filter_str)
def get_optional_models(self, sort_it:bool=False, filter_str:tuple=())->list:
return self._model_prep.prep_models_list(self._optional_models, sort_it, filter_str)
def _set_llm_client(self, url:str, request_type:RequestMode=RequestMode.OPENSOURCE)-> bool:
if not self.is_lm_server_up() or not url:
self._lm_client = None
self._lm_url = url
self._lm_models = None
self.j_mngr.log_events("Local LLM server is not running; aborting client setup.",
TroubleSgltn.Severity.WARNING,
True)
return False
lm_object = OpenAI
key = "No key necessary" #Default value used in LLM front-ends that don't require a key
#Use the requested API
if request_type in (RequestMode.OOBABOOGA, RequestMode.OPENSOURCE):
if not self._lm_key:
self.j_mngr.log_events("Setting Openai client with URL, no key.",
is_trouble=True)
else:
key = self._lm_key
self.j_mngr.log_events("Setting Openai client with URL and key.",
is_trouble=True)
elif request_type == RequestMode.GROQ:
if not self._groq_key:
self.j_mngr.log_events("Attempting to connect to Groq with no key",
TroubleSgltn.Severity.ERROR,
True)
else:
key = self._groq_key
self.j_mngr.log_events("Setting Openai client with URL and Groq key.",
is_trouble=True)
try:
lm_client = lm_object(base_url=url, api_key=key)
self._lm_url = url
self._lm_client = lm_client
except Exception as e:
self.j_mngr.log_events(f"Unable to create LLM client object using URL. Unable to communicate with LLM: {e}",
TroubleSgltn.Severity.ERROR,
True)
return False
return True
@property
def lm_client(self):
return self._lm_client
@property
def lm_url(self):
return self._lm_url
def write_url(self, url:str) -> bool:
# Save the current open source url for startup retrieval of models
url_result = False
if url and url != self._written_url:
url_file = self.j_mngr.append_filename_to_path(self.j_mngr.script_dir, 'OpenSourceURL.txt')
url_result = self.j_mngr.write_string_to_file(url, url_file)
self._written_url = url
self.j_mngr.log_events("Open source LLM URL saved to file.",
TroubleSgltn.Severity.INFO,
True)
return url_result
@lm_url.setter
def lm_url(self, url: str):
if url != self._lm_url or not self._lm_client: # Check if the new URL is different to avoid unnecessary operations
self._lm_url = url
# Reset client and models only if a new URL is provided
self._lm_client = None
#self._lm_models = []
if url: # If the new URL is not empty, update the client
self._set_llm_client(url, self._lm_request_mode)
def is_lm_server_up(self): #should be util in api_requests.py
session = requests.Session()
retries = Retry(total=2, backoff_factor=0, status_forcelist=[500, 502, 503, 504])
session.mount('http://', HTTPAdapter(max_retries=retries))
try:
response = session.head(self._lm_url, timeout=4) # Use HEAD to minimize data transfer
if 200 <= response.status_code <= 300:
self.write_url(self._lm_url) #Save url to a text file
self.j_mngr.log_events(f"Local LLM Server is running with status code: {response.status_code}",
TroubleSgltn.Severity.INFO,
True)
return True
else:
self.write_url(self._lm_url) #Save url to a text file
self.j_mngr.log_events(f"Server returned response code: {response.status_code}",
TroubleSgltn.Severity.INFO,
True)
return True
except requests.RequestException as e:
self.j_mngr.log_events(f"Local LLM Server is not running: {e}",
TroubleSgltn.Severity.WARNING,
True)
return False
@property
def use_examples(self)->bool:
return self._use_examples
@use_examples.setter
def use_examples(self, use_examples: bool):
#Write, sets internal flag
self._use_examples = use_examples
@property
def lm_request_mode(self)->RequestMode:
return self._lm_request_mode
@lm_request_mode.setter
def lm_request_mode(self, mode:RequestMode)-> None:
self._lm_request_mode = mode
@property
def key(self)-> str:
return self._fig_key
@property
def lm_key(self)-> str:
return self._lm_key
@property
def groq_key(self)->str:
return self._groq_key
@property
def anthropic_key(self)->str:
return self._claude_key
@property
def gemini_key(self)->str:
return self._gemini_key
@property
def instruction(self):
return self.figInstruction
@property
def example(self):
if self._use_examples:
return self.figExample
return ""
@property
def example2(self):
if self._use_examples:
return self.figExample2
return ""
@property
def n_Example(self):
if self._use_examples:
return self.fig_n_example
return ""
@property
def n_example2(self):
if self._use_examples:
return self.fig_n_example2
return ""
@property
def style(self):
#make sure the designated default value is present in the list
if "Photograph" not in self.figStyle:
if not isinstance(self.figStyle, list):
self.figStyle = []
self.figStyle.append("Photograph")
return self.figStyle
@property
def ImgInstruction(self):
return self.figImgInstruction
@property
def ImgPromptInstruction(self):
return self.figImgPromptInstruction
@property
def n_Instruction(self):
return self.fig_n_Instruction
@property
def n_ImgPromptInstruction(self):
return self.fig_n_ImgPromptInstruction
@property
def n_ImgInstruction(self):
return self.fig_n_ImgInstruction
@property
def pyexiv2(self)-> Optional[object]:
return self._pyexiv2
@property
def anthropic_client(self)->Optional[object]:
if self._claude_key:
return self._anthropic_client
return None
@property
def openaiClient(self)-> Optional[object]:
if self._fig_key:
return self.figOAIClient
return None
class AI_Chooser:
def __init__(self):
#instantiate Configuration and Help data classes
self.cFig = cFigSingleton()
self.help_data = helpSgltn()
self.j_mngr = json_manager()
self.trbl = TroubleSgltn()
@staticmethod
def select_request_mode(user_selection:str) -> RequestMode:
mode_map = {
"ChatGPT": RequestMode.OPENAI,
"Groq": RequestMode.GROQ,
"Anthropic": RequestMode.CLAUDE,
"LM_Studio": RequestMode.LMSTUDIO,
"Local app (URL)": RequestMode.OPENSOURCE,
"OpenAI compatible http POST": RequestMode.OPENSOURCE,
"http POST Simplified Data": RequestMode.OSSIMPLE,
"Oobabooga API-URL": RequestMode.OOBABOOGA
}
return mode_map.get(user_selection)
@classmethod
def INPUT_TYPES(cls):
cFig=cFigSingleton()
gptfilter = ("gpt","o1")
#Floats have a problem, they go over the max value even when round and step are set, and the node fails. So I set max a little over the expected input value
return {
"required": {
"AI_Service": (["ChatGPT", "Groq", "Anthropic"], {"default": "ChatGPT"}),
"ChatGPT_model": (cFig.get_chat_models(True,gptfilter), {"default": ""}),
"Groq_model": (cFig.get_groq_models(True), {"default": ""}),
"Anthropic_model": (cFig.get_claude_models(True), {"default": ""}),
},
"hidden": {
"unique_id": "UNIQUE_ID",
}
}
RETURN_TYPES = ("DICTIONARY",)
RETURN_NAMES = ("AI_Selection",)
FUNCTION = "gogo"
OUTPUT_NODE = False
CATEGORY = "Plush/Prompt"
def gogo(self, unique_id, AI_Service, ChatGPT_model, Groq_model, Anthropic_model):
ai_dict = {"service": AI_Chooser.select_request_mode(AI_Service), "model": None}
if ai_dict['service'] == RequestMode.OPENAI and ChatGPT_model != "none":
ai_dict['model'] = ChatGPT_model
elif ai_dict['service'] == RequestMode.GROQ and Groq_model != "none":
ai_dict['model'] = Groq_model
elif ai_dict['service'] == RequestMode.CLAUDE and Anthropic_model != "none":
ai_dict['model'] = Anthropic_model
return (ai_dict,)
class Enhancer:
#Build a creative prompt using a ChatGPT model
def __init__(self):
#instantiate Configuration and Help data classes
self.cFig = cFigSingleton()
self.help_data = helpSgltn()
self.j_mngr = json_manager()
self.trbl = TroubleSgltn()
self.ctx = rqst.request_context()
def build_instruction(self, mode, style, prompt_style, elements, artist):
#build the instruction from user input
instruc = ""
if prompt_style == "Narrative":
if mode == InputMode.PROMPT_ONLY:
if self.cFig.n_Instruction:
instruc = self.cFig.n_Instruction
elif mode == InputMode.IMAGE_ONLY:
if self.cFig.n_ImgInstruction:
instruc = self.cFig.n_ImgInstruction
elif mode == InputMode.IMAGE_PROMPT:
if self.cFig.n_ImgPromptInstruction:
instruc = self.cFig.n_ImgPromptInstruction
else: #Prompt_style is Tags
if mode == InputMode.PROMPT_ONLY:
if self.cFig.instruction:
instruc = self.cFig.instruction
elif mode == InputMode.IMAGE_ONLY:
if self.cFig.ImgInstruction:
instruc = self.cFig.ImgInstruction
elif mode == InputMode.IMAGE_PROMPT:
if self.cFig.ImgPromptInstruction:
instruc = self.cFig.ImgPromptInstruction
if instruc.count("{}") >= 2:
instruc = instruc.format(style, elements)
elif instruc.count("{}") == 1:
instruc = instruc.format(style)
if artist >= 1:
art_instruc = " Include {} artist(s) who works in the specifed artistic style by placing the artists' name(s) at the end of the sentence prefaced by 'style of'."
instruc += art_instruc.format(str(artist))
return instruc
def translateModelName(self, model: str)-> str:
#Translate friendly model names to working model names
#Not in use right now, but new models typically go through a period where there's
#no pointer value for the latest models.
if model == "gpt-4 Turbo":
model = "gpt-4-1106-preview"
return model
@staticmethod
def undefined_to_none( sus_var):
"""
Converts the string "undefined" to a None.
Note: ComfyUI returns unconnected UI elements as "undefined"
which causes problems when the node expects these to be handled as falsey
Args:
sus_var(any): The variable that might containt "undefined"
Returns:
None if the variable is set to the string "undefined" or unchanged (any) if not.
"""
return None if sus_var == "undefined" else sus_var
@classmethod
def INPUT_TYPES(cls):
cFig=cFigSingleton()
#Floats have a problem, they go over the max value even when round and step are set, and the node fails. So I set max a little over the expected input value
return {
"required": {
#"GPTmodel": (cFig.get_chat_models(True, 'gpt'),{"default": ""} ),
"creative_latitude" : ("FLOAT", {"max": 1.201, "min": 0.1, "step": 0.1, "display": "number", "round": 0.1, "default": 0.7}),
"tokens" : ("INT", {"max": 8000, "min": 20, "step": 10, "default": 500, "display": "number"}),
"style": (cFig.style,{"default": "Photograph"}),
"artist" : ("INT", {"max": 3, "min": 0, "step": 1, "default": 1, "display": "number"}),
"prompt_style": (["Tags", "Narrative"],{"default": "Tags"}),
"max_elements" : ("INT", {"max": 25, "min": 3, "step": 1, "default": 10, "display": "number"}),
"style_info" : ("BOOLEAN", {"default": False})
},
"hidden": {
"unique_id": "UNIQUE_ID",
},
"optional": {
"AI_Selection":("DICTIONARY", {"default": None}),
"prompt": ("STRING",{"multiline": True, "default": ""}),
"image" : ("IMAGE", {"default": None})
}
}
RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING","STRING")
RETURN_NAMES = ("AI_prompt", "AI_instruction","Style Info", "Help","troubleshooting")
FUNCTION = "gogo"
OUTPUT_NODE = False
CATEGORY = "Plush/Prompt"
def gogo(self, creative_latitude, tokens, style, artist, prompt_style, max_elements, style_info, AI_Selection=None, prompt="", image=None, unique_id=None):
if unique_id:
self.trbl.reset('Style Prompt, Node #'+unique_id)
else:
self.trbl.reset('Style Prompt')
_help = self.help_data.style_prompt_help
CGPT_prompt = ""
instruction = ""
CGPT_styleInfo = ""
if AI_Selection:
ais_model = AI_Selection['model']
else:
self.j_mngr.log_events("You must connect the Plush AI_Chooser to the AI_Selection Input and choose an AI_Service and model to use",
TroubleSgltn.Severity.ERROR,
True)
CGPT_prompt = "Plush AI_Chooser not connected to AI_Selection input, or missing input values"
return(CGPT_prompt, instruction, CGPT_styleInfo, _help, self.trbl.get_troubles())
# unconnected UI elements get passed in as the string "undefined" by ComfyUI
image = self.undefined_to_none(image)
prompt = self.undefined_to_none(prompt)
#Translate any friendly model names
#Convert PyTorch.tensor to B64encoded image
if isinstance(image, torch.Tensor):
image = DalleImage.tensor_to_base64(image)
#build instruction based on user input
mode = 0
if image and prompt:
mode = InputMode.IMAGE_PROMPT
elif image:
mode = InputMode.IMAGE_ONLY
elif prompt:
mode = InputMode.PROMPT_ONLY
instruction = self.build_instruction(mode, style, prompt_style, max_elements, artist)
self.cFig.lm_request_mode = AI_Selection['service']
if AI_Selection['service'] == RequestMode.OPENAI:
self.ctx.request = rqst.oai_object_request()
elif AI_Selection['service'] == RequestMode.GROQ:
self.ctx.request = rqst.oai_object_request()
# set the url so the function making the request will have a properly initialized object.
self.cFig.lm_url = "https://api.groq.com/openai/v1" # Ugh! I've embedded a 'magic value' URL here for the OPENAI API Object because the GROQ API object looks flakey...
elif AI_Selection['service'] == RequestMode.CLAUDE:
self.ctx.request = rqst.claude_request()
if style_info:
self.trbl.set_process_header("Art Style Info:")
#User has request information about the art style. GPT will provide it
sty_prompt = f"Give an 150 word backgrounder on the art style: {style}. Starting with describing what it is, include information about its history and which artists represent the style."
kwargs = { "model": ais_model,
"creative_latitude": creative_latitude,
"tokens": tokens,
"prompt": sty_prompt,
}
CGPT_styleInfo = self.ctx.execute_request(**kwargs)
self.trbl.pop_header()
kwargs = { "model": ais_model,
"creative_latitude": creative_latitude,
"tokens": tokens,
"prompt": prompt,
"instruction": instruction,
"image": image,
}
CGPT_prompt = self.ctx.execute_request(**kwargs)
return (CGPT_prompt, instruction, CGPT_styleInfo, _help, self.trbl.get_troubles())
class AdvPromptEnhancer:
#Advance Prompt Enhancer: User entered Instruction, Prompt and Examples
def __init__(self)-> None:
self.cFig = cFigSingleton()
self.help_data = helpSgltn()
self.j_mngr = json_manager()
self.trbl = TroubleSgltn()
self.ctx = rqst.request_context()
def get_model(self, GPT_model, Groq_model, Anthropic_model, Ollamm_model, Optional_model, connection_type)->str:
if connection_type == "ChatGPT":
return GPT_model
if connection_type == "Groq":
return Groq_model
if connection_type == "Anthropic":
return Anthropic_model
if "Ollama" in connection_type and Ollamm_model != "none":
return Ollamm_model
if Optional_model and Optional_model != "none":
template = {"content": None}
model_dlist = []
model_dlist = self.j_mngr.insert_string_dict(Optional_model,template,"content","::")
if len(model_dlist) > 1:
return model_dlist[1]['content']
return model_dlist[0]['content']
return "none"
@classmethod
def INPUT_TYPES(cls):
cFig = cFigSingleton()
gptfilter = ("gpt","o1")
#open source models are too problematic to implement right now in an environment where you
#can't be sure if the local host server (open source) will be running, and where you can't
#refresh the ui after the initial load.
return {
"required": {
"AI_service": (["ChatGPT", "Groq", "Anthropic", "LM_Studio (URL)", "Ollama (URL)","Local app (URL)", "OpenAI compatible http POST (URL)", "http POST Simplified Data (URL)", "Oobabooga API (URL)"], {"default": "Groq"}),
"ChatGPT_model": (cFig.get_chat_models(True,gptfilter), {"default": ""}),
"Groq_model": (cFig.get_groq_models(True), {"default": ""}),
"Anthropic_model": (cFig.get_claude_models(True), {"default": ""}),
"Ollama_model": (cFig.get_ollama_models(True), {"default": ""}),
"Optional_model": (cFig.get_optional_models(True), {"default": ""}),
"creative_latitude" : ("FLOAT", {"max": 1.901, "min": 0.1, "step": 0.1, "display": "number", "round": 0.1, "default": 0.7}),
"tokens" : ("INT", {"max": 8000, "min": 20, "step": 10, "default": 500, "display": "number"}),
"seed": ("INT", {"default": 9, "min": 0, "max": 0xffffffffffffffff}),
"examples_delimiter":(["Pipe |", "Two newlines", "Two colons ::"], {"default": "Two newlines"}),
"LLM_URL": ("STRING",{"default": cFig.lm_url})
},
"hidden": {
"unique_id": "UNIQUE_ID",
},
"optional": {
"Instruction": ("STRING",{"multiline": True, "default": "", "forceInput": True}),
"Examples_or_Context": ("STRING",{"multiline": True, "default": "", "forceInput": True}),
"Prompt": ("STRING",{"multiline": True, "default": "", "forceInput": True}),
"image" : ("IMAGE", {"default": None})
}
}
RETURN_TYPES = ("STRING", "STRING", "STRING")
RETURN_NAMES = ("LLMprompt", "Help","Troubleshooting")
FUNCTION = "gogo"
OUTPUT_NODE = False
CATEGORY = "Plush/Prompt"
def gogo(self, AI_service, ChatGPT_model, Groq_model, Anthropic_model, Ollama_model, Optional_model, creative_latitude, tokens, seed, examples_delimiter,
LLM_URL:str="", Instruction:str="", Prompt:str = "", Examples_or_Context:str ="",image=None, unique_id=None):
if unique_id:
self.trbl.reset("Advanced Prompt Enhancer, Node #"+unique_id)
else:
self.trbl.reset("Advanced Prompt Enhancer")
_help = self.help_data.adv_prompt_help
# set the value of unconnected inputs to None
Instruction = Enhancer.undefined_to_none(Instruction)
Prompt = Enhancer.undefined_to_none(Prompt)
Examples = Enhancer.undefined_to_none(Examples_or_Context)
LLM_URL = Enhancer.undefined_to_none(LLM_URL)
image = Enhancer.undefined_to_none(image)
remote_model = self.get_model(ChatGPT_model, Groq_model, Anthropic_model, Ollama_model, Optional_model, AI_service)
if remote_model == "none":
self.j_mngr.log_events("No model selected. If you're using a local desktop application, most will just use the loaded model.",
TroubleSgltn.Severity.INFO,
True)
llm_result = "Unable to process request. Make sure the local Open Source Server is running, and you've provided a valid URL. If you're using a remote service (e.g.: ChaTGPT, Groq) make sure your key is valid, and a model is selected"
#Convert PyTorch.tensor to B64encoded image
if isinstance(image, torch.Tensor):
image = DalleImage.tensor_to_base64(image)
#Create a list of dictionaries out of the user provided Examples_or_Context
example_list = []
if Examples:
delimiter = None
if examples_delimiter == "Two newlines":
delimiter = "\n\n"
elif examples_delimiter == "Two colons ::":
delimiter = "::"
elif examples_delimiter == "Pipe |":
delimiter = "|"
example_list = self.j_mngr.build_context(Examples, delimiter)
kwargs = { "model": remote_model,
"creative_latitude": creative_latitude,
"tokens": tokens,
"seed": seed,
"prompt": Prompt,
"instruction": Instruction,
"url": LLM_URL,
"image": image,
"example_list": example_list,
}
if AI_service == 'Local app (URL)' or AI_service == "Groq" or AI_service == "Ollama (URL)":
if AI_service == 'Local app (URL)':
self.cFig.lm_request_mode = RequestMode.OPENSOURCE
elif AI_service == "Groq":
self.cFig.lm_request_mode = RequestMode.GROQ
LLM_URL = "https://api.groq.com/openai/v1" # Ugh! I've embedded a 'magic value' URL here for the OPENAI API Object because the GROQ API object looks flakey...
elif AI_service == "Ollama (URL)":
self.cFig.lm_request_mode = RequestMode.OLLAMA
if not LLM_URL:
self.j_mngr.log_events("'Local app (URL)' specified, but no URL provided or URL is invalid. Enter a valid URL",
TroubleSgltn.Severity.WARNING,
True)
return(llm_result, _help, self.trbl.get_troubles())
# set the url so the function making the request will have a properly initialized object.
self.cFig.lm_url = LLM_URL
if not self.cFig.lm_client:
self.j_mngr.log_events("Open Source LLM server is not running. Aborting request.",
TroubleSgltn.Severity.WARNING,
True)
return(llm_result, _help, self.trbl.get_troubles())
llm_result = ""
self.ctx.request = rqst.oai_object_request( )
llm_result = self.ctx.execute_request(**kwargs)
return(llm_result, _help, self.trbl.get_troubles())
if AI_service == 'Anthropic':
self.cFig.lm_request_mode = RequestMode.CLAUDE
claude_result = ""
self.ctx.request = rqst.claude_request()
claude_result = self.ctx.execute_request(**kwargs)
return(claude_result, _help, self.trbl.get_troubles())
if AI_service == "OpenAI compatible http POST (URL)" or AI_service == "LM_Studio (URL)":
if not LLM_URL:
self.j_mngr.log_events("'OpenAI compatible http POST' specified, but no URL provided or URL is invalid. Enter a valid URL",
TroubleSgltn.Severity.WARNING,
True)
return(llm_result, _help, self.trbl.get_troubles())
self.ctx.request = rqst.oai_web_request()
if AI_service == "LM_Studio (URL)":
self.cFig.lm_request_mode = RequestMode.LMSTUDIO
else:
self.cFig.lm_request_mode = RequestMode.OPENSOURCE
llm_result = self.ctx.execute_request(**kwargs)
return(llm_result, _help, self.trbl.get_troubles())
if AI_service == "http POST Simplified Data (URL)":
if not LLM_URL:
self.j_mngr.log_events("'http POST Simplified Data' specified, but no URL provided or URL is invalid. Enter a valid URL",
TroubleSgltn.Severity.WARNING,
True)
return(llm_result, _help, self.trbl.get_troubles())
self.ctx.request = rqst.oai_web_request()
self.cFig.lm_request_mode = RequestMode.OSSIMPLE
llm_result = self.ctx.execute_request(**kwargs)
return(llm_result, _help, self.trbl.get_troubles())
#Oobabooga via POST
if AI_service == "Oobabooga API (URL)":
if not LLM_URL:
self.j_mngr.log_events("'Oobabooga API-URL' specified, but no URL provided or URL is invalid. Enter a valid URL",
TroubleSgltn.Severity.WARNING,
True)
return(llm_result, _help, self.trbl.get_troubles())
self.ctx.request = rqst.ooba_web_request()
self.cFig.lm_request_mode = RequestMode.OOBABOOGA
llm_result = self.ctx.execute_request(**kwargs)
return(llm_result, _help, self.trbl.get_troubles())
#OpenAI ChatGPT request
self.ctx.request = rqst.oai_object_request()
self.cFig.lm_request_mode = RequestMode.OPENAI
llm_result = self.ctx.execute_request(**kwargs)
return(llm_result, _help, self.trbl.get_troubles())
class DalleImage:
#Accept a user prompt and parameters to produce a Dall_e generated image
def __init__(self):
self.cFig = cFigSingleton()
self.help_data = helpSgltn()
self.j_mngr = json_manager()
self.trbl = TroubleSgltn()
self.ctx = rqst.request_context()
@staticmethod
def b64_to_tensor( b64_image: str) -> tuple[torch.Tensor,torch.Tensor]:
"""
Converts a base64-encoded image to a torch.Tensor.
Note: ComfyUI expects the image tensor in the [N, H, W, C] format.
For example with the shape torch.Size([1, 1024, 1024, 3])
Args:
b64_image (str): The b64 image to convert.
Returns:
torch.Tensor: an image Tensor.
"""
j_mngr = json_manager()
j_mngr.log_events("Converting b64 Image to Torch Tensor Image file",
is_trouble=True)
# Decode the base64 string
image_data = base64.b64decode(b64_image)
# Open the image with PIL and handle EXIF orientation
image = Image.open(BytesIO(image_data))
image = ImageOps.exif_transpose(image)
# Convert to RGBA for potential alpha channel handling
# Dalle doesn't provide an alpha channel, but this is here for
# broad compatibility
image = image.convert("RGBA")
image_np = np.array(image).astype(np.float32) / 255.0 # Normalize
# Split the image into RGB and Alpha channels
rgb_np, alpha_np = image_np[..., :3], image_np[..., 3]
# Convert RGB to PyTorch tensor and ensure it's in the [N, H, W, C] format
tensor_image = torch.from_numpy(rgb_np).unsqueeze(0) # Adds N dimension
# Create mask based on the presence or absence of an alpha channel
if image.mode == 'RGBA':
mask = torch.from_numpy(alpha_np).unsqueeze(0).unsqueeze(0) # Adds N and C dimensions
else: # Fallback if no alpha channel is present
mask = torch.zeros((1, tensor_image.shape[2], tensor_image.shape[3]), dtype=torch.float32) # [N, H, W]
return tensor_image, mask
@staticmethod
def tensor_to_base64(tensor: torch.Tensor) -> str:
"""
Converts a PyTorch tensor to a base64-encoded image.
Note: ComfyUI provides the image tensor in the [N, H, W, C] format.
For example with the shape torch.Size([1, 1024, 1024, 3])
Args:
tensor (torch.Tensor): The image tensor to convert.
Returns:
str: Base64-encoded image string.
"""
j_mngr = json_manager()
j_mngr.log_events("Converting Torch Tensor image to b64 Image file",
is_trouble=True)
# Convert tensor to PIL Image
if tensor.ndim == 4:
tensor = tensor.squeeze(0) # Remove batch dimension if present
pil_image = Image.fromarray((tensor.numpy() * 255).astype('uint8'))
# Save PIL Image to a buffer
buffer = BytesIO()
pil_image.save(buffer, format="PNG") # Can change to JPEG if preferred
buffer.seek(0)
# Encode buffer to base64
base64_image = base64.b64encode(buffer.read()).decode('utf-8')
return base64_image
@staticmethod
def tensor_to_bytes(tensor: torch.Tensor) -> BytesIO:
"""
Converts a PyTorch tensor to a bytes object.
Args:
tensor (torch.Tensor): The image tensor to convert.
Returns:
BytesIO: BytesIO object containing the image data.
"""
# Convert tensor to PIL Image
if tensor.ndim == 4:
tensor = tensor.squeeze(0) # Remove batch dimension if present
pil_image = Image.fromarray((tensor.numpy() * 255).astype('uint8'))
# Save PIL Image to a buffer
buffer = BytesIO()
pil_image.save(buffer, format="PNG") # Can change to JPEG if preferred
buffer.seek(0)
return buffer
@classmethod
def INPUT_TYPES(cls):
#dall-e-2 API requires differnt input parameters as compared to dall-e-3, at this point I'll just use dall-e-3
# "batch_size": ("INT", {"max": 8, "min": 1, "step": 1, "default": 1, "display": "number"})
# Possible future implentation of batch_sizes greater than one.
# "image" : ("IMAGE", {"forceInput": True}),
return {
"required": {
"GPTmodel": (["dall-e-3",], ),
"prompt": ("STRING",{"multiline": True, "forceInput": True}),
"image_size": (["1792x1024", "1024x1792", "1024x1024"], {"default": "1024x1024"} ),