scripts.py 40 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025
  1. import os
  2. import re
  3. import sys
  4. import inspect
  5. from collections import namedtuple
  6. from dataclasses import dataclass
  7. import gradio as gr
  8. from modules import shared, paths, script_callbacks, extensions, script_loading, scripts_postprocessing, errors, timer, util
  9. topological_sort = util.topological_sort
  10. AlwaysVisible = object()
  11. class MaskBlendArgs:
  12. def __init__(self, current_latent, nmask, init_latent, mask, blended_latent, denoiser=None, sigma=None):
  13. self.current_latent = current_latent
  14. self.nmask = nmask
  15. self.init_latent = init_latent
  16. self.mask = mask
  17. self.blended_latent = blended_latent
  18. self.denoiser = denoiser
  19. self.is_final_blend = denoiser is None
  20. self.sigma = sigma
  21. class PostSampleArgs:
  22. def __init__(self, samples):
  23. self.samples = samples
  24. class PostprocessImageArgs:
  25. def __init__(self, image):
  26. self.image = image
  27. class PostProcessMaskOverlayArgs:
  28. def __init__(self, index, mask_for_overlay, overlay_image):
  29. self.index = index
  30. self.mask_for_overlay = mask_for_overlay
  31. self.overlay_image = overlay_image
  32. class PostprocessBatchListArgs:
  33. def __init__(self, images):
  34. self.images = images
  35. @dataclass
  36. class OnComponent:
  37. component: gr.blocks.Block
  38. class Script:
  39. name = None
  40. """script's internal name derived from title"""
  41. section = None
  42. """name of UI section that the script's controls will be placed into"""
  43. filename = None
  44. args_from = None
  45. args_to = None
  46. alwayson = False
  47. is_txt2img = False
  48. is_img2img = False
  49. tabname = None
  50. group = None
  51. """A gr.Group component that has all script's UI inside it."""
  52. create_group = True
  53. """If False, for alwayson scripts, a group component will not be created."""
  54. infotext_fields = None
  55. """if set in ui(), this is a list of pairs of gradio component + text; the text will be used when
  56. parsing infotext to set the value for the component; see ui.py's txt2img_paste_fields for an example
  57. """
  58. paste_field_names = None
  59. """if set in ui(), this is a list of names of infotext fields; the fields will be sent through the
  60. various "Send to <X>" buttons when clicked
  61. """
  62. api_info = None
  63. """Generated value of type modules.api.models.ScriptInfo with information about the script for API"""
  64. on_before_component_elem_id = None
  65. """list of callbacks to be called before a component with an elem_id is created"""
  66. on_after_component_elem_id = None
  67. """list of callbacks to be called after a component with an elem_id is created"""
  68. setup_for_ui_only = False
  69. """If true, the script setup will only be run in Gradio UI, not in API"""
  70. controls = None
  71. """A list of controls returned by the ui()."""
  72. def title(self):
  73. """this function should return the title of the script. This is what will be displayed in the dropdown menu."""
  74. raise NotImplementedError()
  75. def ui(self, is_img2img):
  76. """this function should create gradio UI elements. See https://gradio.app/docs/#components
  77. The return value should be an array of all components that are used in processing.
  78. Values of those returned components will be passed to run() and process() functions.
  79. """
  80. pass
  81. def show(self, is_img2img):
  82. """
  83. is_img2img is True if this function is called for the img2img interface, and False otherwise
  84. This function should return:
  85. - False if the script should not be shown in UI at all
  86. - True if the script should be shown in UI if it's selected in the scripts dropdown
  87. - script.AlwaysVisible if the script should be shown in UI at all times
  88. """
  89. return True
  90. def run(self, p, *args):
  91. """
  92. This function is called if the script has been selected in the script dropdown.
  93. It must do all processing and return the Processed object with results, same as
  94. one returned by processing.process_images.
  95. Usually the processing is done by calling the processing.process_images function.
  96. args contains all values returned by components from ui()
  97. """
  98. pass
  99. def setup(self, p, *args):
  100. """For AlwaysVisible scripts, this function is called when the processing object is set up, before any processing starts.
  101. args contains all values returned by components from ui().
  102. """
  103. pass
  104. def before_process(self, p, *args):
  105. """
  106. This function is called very early during processing begins for AlwaysVisible scripts.
  107. You can modify the processing object (p) here, inject hooks, etc.
  108. args contains all values returned by components from ui()
  109. """
  110. pass
  111. def process(self, p, *args):
  112. """
  113. This function is called before processing begins for AlwaysVisible scripts.
  114. You can modify the processing object (p) here, inject hooks, etc.
  115. args contains all values returned by components from ui()
  116. """
  117. pass
  118. def before_process_batch(self, p, *args, **kwargs):
  119. """
  120. Called before extra networks are parsed from the prompt, so you can add
  121. new extra network keywords to the prompt with this callback.
  122. **kwargs will have those items:
  123. - batch_number - index of current batch, from 0 to number of batches-1
  124. - prompts - list of prompts for current batch; you can change contents of this list but changing the number of entries will likely break things
  125. - seeds - list of seeds for current batch
  126. - subseeds - list of subseeds for current batch
  127. """
  128. pass
  129. def after_extra_networks_activate(self, p, *args, **kwargs):
  130. """
  131. Called after extra networks activation, before conds calculation
  132. allow modification of the network after extra networks activation been applied
  133. won't be call if p.disable_extra_networks
  134. **kwargs will have those items:
  135. - batch_number - index of current batch, from 0 to number of batches-1
  136. - prompts - list of prompts for current batch; you can change contents of this list but changing the number of entries will likely break things
  137. - seeds - list of seeds for current batch
  138. - subseeds - list of subseeds for current batch
  139. - extra_network_data - list of ExtraNetworkParams for current stage
  140. """
  141. pass
  142. def process_batch(self, p, *args, **kwargs):
  143. """
  144. Same as process(), but called for every batch.
  145. **kwargs will have those items:
  146. - batch_number - index of current batch, from 0 to number of batches-1
  147. - prompts - list of prompts for current batch; you can change contents of this list but changing the number of entries will likely break things
  148. - seeds - list of seeds for current batch
  149. - subseeds - list of subseeds for current batch
  150. """
  151. pass
  152. def postprocess_batch(self, p, *args, **kwargs):
  153. """
  154. Same as process_batch(), but called for every batch after it has been generated.
  155. **kwargs will have same items as process_batch, and also:
  156. - batch_number - index of current batch, from 0 to number of batches-1
  157. - images - torch tensor with all generated images, with values ranging from 0 to 1;
  158. """
  159. pass
  160. def postprocess_batch_list(self, p, pp: PostprocessBatchListArgs, *args, **kwargs):
  161. """
  162. Same as postprocess_batch(), but receives batch images as a list of 3D tensors instead of a 4D tensor.
  163. This is useful when you want to update the entire batch instead of individual images.
  164. You can modify the postprocessing object (pp) to update the images in the batch, remove images, add images, etc.
  165. If the number of images is different from the batch size when returning,
  166. then the script has the responsibility to also update the following attributes in the processing object (p):
  167. - p.prompts
  168. - p.negative_prompts
  169. - p.seeds
  170. - p.subseeds
  171. **kwargs will have same items as process_batch, and also:
  172. - batch_number - index of current batch, from 0 to number of batches-1
  173. """
  174. pass
  175. def on_mask_blend(self, p, mba: MaskBlendArgs, *args):
  176. """
  177. Called in inpainting mode when the original content is blended with the inpainted content.
  178. This is called at every step in the denoising process and once at the end.
  179. If is_final_blend is true, this is called for the final blending stage.
  180. Otherwise, denoiser and sigma are defined and may be used to inform the procedure.
  181. """
  182. pass
  183. def post_sample(self, p, ps: PostSampleArgs, *args):
  184. """
  185. Called after the samples have been generated,
  186. but before they have been decoded by the VAE, if applicable.
  187. Check getattr(samples, 'already_decoded', False) to test if the images are decoded.
  188. """
  189. pass
  190. def postprocess_image(self, p, pp: PostprocessImageArgs, *args):
  191. """
  192. Called for every image after it has been generated.
  193. """
  194. pass
  195. def postprocess_maskoverlay(self, p, ppmo: PostProcessMaskOverlayArgs, *args):
  196. """
  197. Called for every image after it has been generated.
  198. """
  199. pass
  200. def postprocess_image_after_composite(self, p, pp: PostprocessImageArgs, *args):
  201. """
  202. Called for every image after it has been generated.
  203. Same as postprocess_image but after inpaint_full_res composite
  204. So that it operates on the full image instead of the inpaint_full_res crop region.
  205. """
  206. pass
  207. def postprocess(self, p, processed, *args):
  208. """
  209. This function is called after processing ends for AlwaysVisible scripts.
  210. args contains all values returned by components from ui()
  211. """
  212. pass
  213. def before_component(self, component, **kwargs):
  214. """
  215. Called before a component is created.
  216. Use elem_id/label fields of kwargs to figure out which component it is.
  217. This can be useful to inject your own components somewhere in the middle of vanilla UI.
  218. You can return created components in the ui() function to add them to the list of arguments for your processing functions
  219. """
  220. pass
  221. def after_component(self, component, **kwargs):
  222. """
  223. Called after a component is created. Same as above.
  224. """
  225. pass
  226. def on_before_component(self, callback, *, elem_id):
  227. """
  228. Calls callback before a component is created. The callback function is called with a single argument of type OnComponent.
  229. May be called in show() or ui() - but it may be too late in latter as some components may already be created.
  230. This function is an alternative to before_component in that it also cllows to run before a component is created, but
  231. it doesn't require to be called for every created component - just for the one you need.
  232. """
  233. if self.on_before_component_elem_id is None:
  234. self.on_before_component_elem_id = []
  235. self.on_before_component_elem_id.append((elem_id, callback))
  236. def on_after_component(self, callback, *, elem_id):
  237. """
  238. Calls callback after a component is created. The callback function is called with a single argument of type OnComponent.
  239. """
  240. if self.on_after_component_elem_id is None:
  241. self.on_after_component_elem_id = []
  242. self.on_after_component_elem_id.append((elem_id, callback))
  243. def describe(self):
  244. """unused"""
  245. return ""
  246. def elem_id(self, item_id):
  247. """helper function to generate id for a HTML element, constructs final id out of script name, tab and user-supplied item_id"""
  248. need_tabname = self.show(True) == self.show(False)
  249. tabkind = 'img2img' if self.is_img2img else 'txt2img'
  250. tabname = f"{tabkind}_" if need_tabname else ""
  251. title = re.sub(r'[^a-z_0-9]', '', re.sub(r'\s', '_', self.title().lower()))
  252. return f'script_{tabname}{title}_{item_id}'
  253. def before_hr(self, p, *args):
  254. """
  255. This function is called before hires fix start.
  256. """
  257. pass
  258. class ScriptBuiltinUI(Script):
  259. setup_for_ui_only = True
  260. def elem_id(self, item_id):
  261. """helper function to generate id for a HTML element, constructs final id out of tab and user-supplied item_id"""
  262. need_tabname = self.show(True) == self.show(False)
  263. tabname = ('img2img' if self.is_img2img else 'txt2img') + "_" if need_tabname else ""
  264. return f'{tabname}{item_id}'
  265. def show(self, is_img2img):
  266. return AlwaysVisible
  267. current_basedir = paths.script_path
  268. def basedir():
  269. """returns the base directory for the current script. For scripts in the main scripts directory,
  270. this is the main directory (where webui.py resides), and for scripts in extensions directory
  271. (ie extensions/aesthetic/script/aesthetic.py), this is extension's directory (extensions/aesthetic)
  272. """
  273. return current_basedir
  274. ScriptFile = namedtuple("ScriptFile", ["basedir", "filename", "path"])
  275. scripts_data = []
  276. postprocessing_scripts_data = []
  277. ScriptClassData = namedtuple("ScriptClassData", ["script_class", "path", "basedir", "module"])
  278. @dataclass
  279. class ScriptWithDependencies:
  280. script_canonical_name: str
  281. file: ScriptFile
  282. requires: list
  283. load_before: list
  284. load_after: list
  285. def list_scripts(scriptdirname, extension, *, include_extensions=True):
  286. scripts = {}
  287. loaded_extensions = {ext.canonical_name: ext for ext in extensions.active()}
  288. loaded_extensions_scripts = {ext.canonical_name: [] for ext in extensions.active()}
  289. # build script dependency map
  290. root_script_basedir = os.path.join(paths.script_path, scriptdirname)
  291. if os.path.exists(root_script_basedir):
  292. for filename in sorted(os.listdir(root_script_basedir)):
  293. if not os.path.isfile(os.path.join(root_script_basedir, filename)):
  294. continue
  295. if os.path.splitext(filename)[1].lower() != extension:
  296. continue
  297. script_file = ScriptFile(paths.script_path, filename, os.path.join(root_script_basedir, filename))
  298. scripts[filename] = ScriptWithDependencies(filename, script_file, [], [], [])
  299. if include_extensions:
  300. for ext in extensions.active():
  301. extension_scripts_list = ext.list_files(scriptdirname, extension)
  302. for extension_script in extension_scripts_list:
  303. if not os.path.isfile(extension_script.path):
  304. continue
  305. script_canonical_name = ("builtin/" if ext.is_builtin else "") + ext.canonical_name + "/" + extension_script.filename
  306. relative_path = scriptdirname + "/" + extension_script.filename
  307. script = ScriptWithDependencies(
  308. script_canonical_name=script_canonical_name,
  309. file=extension_script,
  310. requires=ext.metadata.get_script_requirements("Requires", relative_path, scriptdirname),
  311. load_before=ext.metadata.get_script_requirements("Before", relative_path, scriptdirname),
  312. load_after=ext.metadata.get_script_requirements("After", relative_path, scriptdirname),
  313. )
  314. scripts[script_canonical_name] = script
  315. loaded_extensions_scripts[ext.canonical_name].append(script)
  316. for script_canonical_name, script in scripts.items():
  317. # load before requires inverse dependency
  318. # in this case, append the script name into the load_after list of the specified script
  319. for load_before in script.load_before:
  320. # if this requires an individual script to be loaded before
  321. other_script = scripts.get(load_before)
  322. if other_script:
  323. other_script.load_after.append(script_canonical_name)
  324. # if this requires an extension
  325. other_extension_scripts = loaded_extensions_scripts.get(load_before)
  326. if other_extension_scripts:
  327. for other_script in other_extension_scripts:
  328. other_script.load_after.append(script_canonical_name)
  329. # if After mentions an extension, remove it and instead add all of its scripts
  330. for load_after in list(script.load_after):
  331. if load_after not in scripts and load_after in loaded_extensions_scripts:
  332. script.load_after.remove(load_after)
  333. for other_script in loaded_extensions_scripts.get(load_after, []):
  334. script.load_after.append(other_script.script_canonical_name)
  335. dependencies = {}
  336. for script_canonical_name, script in scripts.items():
  337. for required_script in script.requires:
  338. if required_script not in scripts and required_script not in loaded_extensions:
  339. errors.report(f'Script "{script_canonical_name}" requires "{required_script}" to be loaded, but it is not.', exc_info=False)
  340. dependencies[script_canonical_name] = script.load_after
  341. ordered_scripts = topological_sort(dependencies)
  342. scripts_list = [scripts[script_canonical_name].file for script_canonical_name in ordered_scripts]
  343. return scripts_list
  344. def list_files_with_name(filename):
  345. res = []
  346. dirs = [paths.script_path] + [ext.path for ext in extensions.active()]
  347. for dirpath in dirs:
  348. if not os.path.isdir(dirpath):
  349. continue
  350. path = os.path.join(dirpath, filename)
  351. if os.path.isfile(path):
  352. res.append(path)
  353. return res
  354. def load_scripts():
  355. global current_basedir
  356. scripts_data.clear()
  357. postprocessing_scripts_data.clear()
  358. script_callbacks.clear_callbacks()
  359. scripts_list = list_scripts("scripts", ".py") + list_scripts("modules/processing_scripts", ".py", include_extensions=False)
  360. syspath = sys.path
  361. def register_scripts_from_module(module):
  362. for script_class in module.__dict__.values():
  363. if not inspect.isclass(script_class):
  364. continue
  365. if issubclass(script_class, Script):
  366. scripts_data.append(ScriptClassData(script_class, scriptfile.path, scriptfile.basedir, module))
  367. elif issubclass(script_class, scripts_postprocessing.ScriptPostprocessing):
  368. postprocessing_scripts_data.append(ScriptClassData(script_class, scriptfile.path, scriptfile.basedir, module))
  369. # here the scripts_list is already ordered
  370. # processing_script is not considered though
  371. for scriptfile in scripts_list:
  372. try:
  373. if scriptfile.basedir != paths.script_path:
  374. sys.path = [scriptfile.basedir] + sys.path
  375. current_basedir = scriptfile.basedir
  376. script_module = script_loading.load_module(scriptfile.path)
  377. register_scripts_from_module(script_module)
  378. except Exception:
  379. errors.report(f"Error loading script: {scriptfile.filename}", exc_info=True)
  380. finally:
  381. sys.path = syspath
  382. current_basedir = paths.script_path
  383. timer.startup_timer.record(scriptfile.filename)
  384. global scripts_txt2img, scripts_img2img, scripts_postproc
  385. scripts_txt2img = ScriptRunner()
  386. scripts_img2img = ScriptRunner()
  387. scripts_postproc = scripts_postprocessing.ScriptPostprocessingRunner()
  388. def wrap_call(func, filename, funcname, *args, default=None, **kwargs):
  389. try:
  390. return func(*args, **kwargs)
  391. except Exception:
  392. errors.report(f"Error calling: {filename}/{funcname}", exc_info=True)
  393. return default
  394. class ScriptRunner:
  395. def __init__(self):
  396. self.scripts = []
  397. self.selectable_scripts = []
  398. self.alwayson_scripts = []
  399. self.titles = []
  400. self.title_map = {}
  401. self.infotext_fields = []
  402. self.paste_field_names = []
  403. self.inputs = [None]
  404. self.callback_map = {}
  405. self.callback_names = [
  406. 'before_process',
  407. 'process',
  408. 'before_process_batch',
  409. 'after_extra_networks_activate',
  410. 'process_batch',
  411. 'postprocess',
  412. 'postprocess_batch',
  413. 'postprocess_batch_list',
  414. 'post_sample',
  415. 'on_mask_blend',
  416. 'postprocess_image',
  417. 'postprocess_maskoverlay',
  418. 'postprocess_image_after_composite',
  419. 'before_component',
  420. 'after_component',
  421. ]
  422. self.on_before_component_elem_id = {}
  423. """dict of callbacks to be called before an element is created; key=elem_id, value=list of callbacks"""
  424. self.on_after_component_elem_id = {}
  425. """dict of callbacks to be called after an element is created; key=elem_id, value=list of callbacks"""
  426. def initialize_scripts(self, is_img2img):
  427. from modules import scripts_auto_postprocessing
  428. self.scripts.clear()
  429. self.alwayson_scripts.clear()
  430. self.selectable_scripts.clear()
  431. auto_processing_scripts = scripts_auto_postprocessing.create_auto_preprocessing_script_data()
  432. for script_data in auto_processing_scripts + scripts_data:
  433. try:
  434. script = script_data.script_class()
  435. except Exception:
  436. errors.report(f"Error # failed to initialize Script {script_data.module}: ", exc_info=True)
  437. continue
  438. script.filename = script_data.path
  439. script.is_txt2img = not is_img2img
  440. script.is_img2img = is_img2img
  441. script.tabname = "img2img" if is_img2img else "txt2img"
  442. visibility = script.show(script.is_img2img)
  443. if visibility == AlwaysVisible:
  444. self.scripts.append(script)
  445. self.alwayson_scripts.append(script)
  446. script.alwayson = True
  447. elif visibility:
  448. self.scripts.append(script)
  449. self.selectable_scripts.append(script)
  450. self.callback_map.clear()
  451. self.apply_on_before_component_callbacks()
  452. def apply_on_before_component_callbacks(self):
  453. for script in self.scripts:
  454. on_before = script.on_before_component_elem_id or []
  455. on_after = script.on_after_component_elem_id or []
  456. for elem_id, callback in on_before:
  457. if elem_id not in self.on_before_component_elem_id:
  458. self.on_before_component_elem_id[elem_id] = []
  459. self.on_before_component_elem_id[elem_id].append((callback, script))
  460. for elem_id, callback in on_after:
  461. if elem_id not in self.on_after_component_elem_id:
  462. self.on_after_component_elem_id[elem_id] = []
  463. self.on_after_component_elem_id[elem_id].append((callback, script))
  464. on_before.clear()
  465. on_after.clear()
  466. def create_script_ui(self, script):
  467. script.args_from = len(self.inputs)
  468. script.args_to = len(self.inputs)
  469. try:
  470. self.create_script_ui_inner(script)
  471. except Exception:
  472. errors.report(f"Error creating UI for {script.name}: ", exc_info=True)
  473. def create_script_ui_inner(self, script):
  474. import modules.api.models as api_models
  475. controls = wrap_call(script.ui, script.filename, "ui", script.is_img2img)
  476. script.controls = controls
  477. if controls is None:
  478. return
  479. script.name = wrap_call(script.title, script.filename, "title", default=script.filename).lower()
  480. api_args = []
  481. for control in controls:
  482. control.custom_script_source = os.path.basename(script.filename)
  483. arg_info = api_models.ScriptArg(label=control.label or "")
  484. for field in ("value", "minimum", "maximum", "step"):
  485. v = getattr(control, field, None)
  486. if v is not None:
  487. setattr(arg_info, field, v)
  488. choices = getattr(control, 'choices', None) # as of gradio 3.41, some items in choices are strings, and some are tuples where the first elem is the string
  489. if choices is not None:
  490. arg_info.choices = [x[0] if isinstance(x, tuple) else x for x in choices]
  491. api_args.append(arg_info)
  492. script.api_info = api_models.ScriptInfo(
  493. name=script.name,
  494. is_img2img=script.is_img2img,
  495. is_alwayson=script.alwayson,
  496. args=api_args,
  497. )
  498. if script.infotext_fields is not None:
  499. self.infotext_fields += script.infotext_fields
  500. if script.paste_field_names is not None:
  501. self.paste_field_names += script.paste_field_names
  502. self.inputs += controls
  503. script.args_to = len(self.inputs)
  504. def setup_ui_for_section(self, section, scriptlist=None):
  505. if scriptlist is None:
  506. scriptlist = self.alwayson_scripts
  507. for script in scriptlist:
  508. if script.alwayson and script.section != section:
  509. continue
  510. if script.create_group:
  511. with gr.Group(visible=script.alwayson) as group:
  512. self.create_script_ui(script)
  513. script.group = group
  514. else:
  515. self.create_script_ui(script)
  516. def prepare_ui(self):
  517. self.inputs = [None]
  518. def setup_ui(self):
  519. all_titles = [wrap_call(script.title, script.filename, "title") or script.filename for script in self.scripts]
  520. self.title_map = {title.lower(): script for title, script in zip(all_titles, self.scripts)}
  521. self.titles = [wrap_call(script.title, script.filename, "title") or f"{script.filename} [error]" for script in self.selectable_scripts]
  522. self.setup_ui_for_section(None)
  523. dropdown = gr.Dropdown(label="Script", elem_id="script_list", choices=["None"] + self.titles, value="None", type="index")
  524. self.inputs[0] = dropdown
  525. self.setup_ui_for_section(None, self.selectable_scripts)
  526. def select_script(script_index):
  527. if script_index is None:
  528. script_index = 0
  529. selected_script = self.selectable_scripts[script_index - 1] if script_index>0 else None
  530. return [gr.update(visible=selected_script == s) for s in self.selectable_scripts]
  531. def init_field(title):
  532. """called when an initial value is set from ui-config.json to show script's UI components"""
  533. if title == 'None':
  534. return
  535. script_index = self.titles.index(title)
  536. self.selectable_scripts[script_index].group.visible = True
  537. dropdown.init_field = init_field
  538. dropdown.change(
  539. fn=select_script,
  540. inputs=[dropdown],
  541. outputs=[script.group for script in self.selectable_scripts]
  542. )
  543. self.script_load_ctr = 0
  544. def onload_script_visibility(params):
  545. title = params.get('Script', None)
  546. if title:
  547. try:
  548. title_index = self.titles.index(title)
  549. visibility = title_index == self.script_load_ctr
  550. self.script_load_ctr = (self.script_load_ctr + 1) % len(self.titles)
  551. return gr.update(visible=visibility)
  552. except ValueError:
  553. params['Script'] = None
  554. massage = f'Cannot find Script: "{title}"'
  555. print(massage)
  556. gr.Warning(massage)
  557. return gr.update(visible=False)
  558. self.infotext_fields.append((dropdown, lambda x: gr.update(value=x.get('Script', 'None'))))
  559. self.infotext_fields.extend([(script.group, onload_script_visibility) for script in self.selectable_scripts])
  560. self.apply_on_before_component_callbacks()
  561. return self.inputs
  562. def run(self, p, *args):
  563. script_index = args[0]
  564. if script_index == 0 or script_index is None:
  565. return None
  566. script = self.selectable_scripts[script_index-1]
  567. if script is None:
  568. return None
  569. script_args = args[script.args_from:script.args_to]
  570. processed = script.run(p, *script_args)
  571. shared.total_tqdm.clear()
  572. return processed
  573. def list_scripts_for_method(self, method_name):
  574. if method_name in ('before_component', 'after_component'):
  575. return self.scripts
  576. else:
  577. return self.alwayson_scripts
  578. def create_ordered_callbacks_list(self, method_name, *, enable_user_sort=True):
  579. script_list = self.list_scripts_for_method(method_name)
  580. category = f'script_{method_name}'
  581. callbacks = []
  582. for script in script_list:
  583. if getattr(script.__class__, method_name, None) == getattr(Script, method_name, None):
  584. continue
  585. script_callbacks.add_callback(callbacks, script, category=category, name=script.__class__.__name__, filename=script.filename)
  586. return script_callbacks.sort_callbacks(category, callbacks, enable_user_sort=enable_user_sort)
  587. def ordered_callbacks(self, method_name, *, enable_user_sort=True):
  588. script_list = self.list_scripts_for_method(method_name)
  589. category = f'script_{method_name}'
  590. scrpts_len, callbacks = self.callback_map.get(category, (-1, None))
  591. if callbacks is None or scrpts_len != len(script_list):
  592. callbacks = self.create_ordered_callbacks_list(method_name, enable_user_sort=enable_user_sort)
  593. self.callback_map[category] = len(script_list), callbacks
  594. return callbacks
  595. def ordered_scripts(self, method_name):
  596. return [x.callback for x in self.ordered_callbacks(method_name)]
  597. def before_process(self, p):
  598. for script in self.ordered_scripts('before_process'):
  599. try:
  600. script_args = p.script_args[script.args_from:script.args_to]
  601. script.before_process(p, *script_args)
  602. except Exception:
  603. errors.report(f"Error running before_process: {script.filename}", exc_info=True)
  604. def process(self, p):
  605. for script in self.ordered_scripts('process'):
  606. try:
  607. script_args = p.script_args[script.args_from:script.args_to]
  608. script.process(p, *script_args)
  609. except Exception:
  610. errors.report(f"Error running process: {script.filename}", exc_info=True)
  611. def before_process_batch(self, p, **kwargs):
  612. for script in self.ordered_scripts('before_process_batch'):
  613. try:
  614. script_args = p.script_args[script.args_from:script.args_to]
  615. script.before_process_batch(p, *script_args, **kwargs)
  616. except Exception:
  617. errors.report(f"Error running before_process_batch: {script.filename}", exc_info=True)
  618. def after_extra_networks_activate(self, p, **kwargs):
  619. for script in self.ordered_scripts('after_extra_networks_activate'):
  620. try:
  621. script_args = p.script_args[script.args_from:script.args_to]
  622. script.after_extra_networks_activate(p, *script_args, **kwargs)
  623. except Exception:
  624. errors.report(f"Error running after_extra_networks_activate: {script.filename}", exc_info=True)
  625. def process_batch(self, p, **kwargs):
  626. for script in self.ordered_scripts('process_batch'):
  627. try:
  628. script_args = p.script_args[script.args_from:script.args_to]
  629. script.process_batch(p, *script_args, **kwargs)
  630. except Exception:
  631. errors.report(f"Error running process_batch: {script.filename}", exc_info=True)
  632. def postprocess(self, p, processed):
  633. for script in self.ordered_scripts('postprocess'):
  634. try:
  635. script_args = p.script_args[script.args_from:script.args_to]
  636. script.postprocess(p, processed, *script_args)
  637. except Exception:
  638. errors.report(f"Error running postprocess: {script.filename}", exc_info=True)
  639. def postprocess_batch(self, p, images, **kwargs):
  640. for script in self.ordered_scripts('postprocess_batch'):
  641. try:
  642. script_args = p.script_args[script.args_from:script.args_to]
  643. script.postprocess_batch(p, *script_args, images=images, **kwargs)
  644. except Exception:
  645. errors.report(f"Error running postprocess_batch: {script.filename}", exc_info=True)
  646. def postprocess_batch_list(self, p, pp: PostprocessBatchListArgs, **kwargs):
  647. for script in self.ordered_scripts('postprocess_batch_list'):
  648. try:
  649. script_args = p.script_args[script.args_from:script.args_to]
  650. script.postprocess_batch_list(p, pp, *script_args, **kwargs)
  651. except Exception:
  652. errors.report(f"Error running postprocess_batch_list: {script.filename}", exc_info=True)
  653. def post_sample(self, p, ps: PostSampleArgs):
  654. for script in self.ordered_scripts('post_sample'):
  655. try:
  656. script_args = p.script_args[script.args_from:script.args_to]
  657. script.post_sample(p, ps, *script_args)
  658. except Exception:
  659. errors.report(f"Error running post_sample: {script.filename}", exc_info=True)
  660. def on_mask_blend(self, p, mba: MaskBlendArgs):
  661. for script in self.ordered_scripts('on_mask_blend'):
  662. try:
  663. script_args = p.script_args[script.args_from:script.args_to]
  664. script.on_mask_blend(p, mba, *script_args)
  665. except Exception:
  666. errors.report(f"Error running post_sample: {script.filename}", exc_info=True)
  667. def postprocess_image(self, p, pp: PostprocessImageArgs):
  668. for script in self.ordered_scripts('postprocess_image'):
  669. try:
  670. script_args = p.script_args[script.args_from:script.args_to]
  671. script.postprocess_image(p, pp, *script_args)
  672. except Exception:
  673. errors.report(f"Error running postprocess_image: {script.filename}", exc_info=True)
  674. def postprocess_maskoverlay(self, p, ppmo: PostProcessMaskOverlayArgs):
  675. for script in self.ordered_scripts('postprocess_maskoverlay'):
  676. try:
  677. script_args = p.script_args[script.args_from:script.args_to]
  678. script.postprocess_maskoverlay(p, ppmo, *script_args)
  679. except Exception:
  680. errors.report(f"Error running postprocess_image: {script.filename}", exc_info=True)
  681. def postprocess_image_after_composite(self, p, pp: PostprocessImageArgs):
  682. for script in self.ordered_scripts('postprocess_image_after_composite'):
  683. try:
  684. script_args = p.script_args[script.args_from:script.args_to]
  685. script.postprocess_image_after_composite(p, pp, *script_args)
  686. except Exception:
  687. errors.report(f"Error running postprocess_image_after_composite: {script.filename}", exc_info=True)
  688. def before_component(self, component, **kwargs):
  689. for callback, script in self.on_before_component_elem_id.get(kwargs.get("elem_id"), []):
  690. try:
  691. callback(OnComponent(component=component))
  692. except Exception:
  693. errors.report(f"Error running on_before_component: {script.filename}", exc_info=True)
  694. for script in self.ordered_scripts('before_component'):
  695. try:
  696. script.before_component(component, **kwargs)
  697. except Exception:
  698. errors.report(f"Error running before_component: {script.filename}", exc_info=True)
  699. def after_component(self, component, **kwargs):
  700. for callback, script in self.on_after_component_elem_id.get(component.elem_id, []):
  701. try:
  702. callback(OnComponent(component=component))
  703. except Exception:
  704. errors.report(f"Error running on_after_component: {script.filename}", exc_info=True)
  705. for script in self.ordered_scripts('after_component'):
  706. try:
  707. script.after_component(component, **kwargs)
  708. except Exception:
  709. errors.report(f"Error running after_component: {script.filename}", exc_info=True)
  710. def script(self, title):
  711. return self.title_map.get(title.lower())
  712. def reload_sources(self, cache):
  713. for si, script in list(enumerate(self.scripts)):
  714. args_from = script.args_from
  715. args_to = script.args_to
  716. filename = script.filename
  717. module = cache.get(filename, None)
  718. if module is None:
  719. module = script_loading.load_module(script.filename)
  720. cache[filename] = module
  721. for script_class in module.__dict__.values():
  722. if type(script_class) == type and issubclass(script_class, Script):
  723. self.scripts[si] = script_class()
  724. self.scripts[si].filename = filename
  725. self.scripts[si].args_from = args_from
  726. self.scripts[si].args_to = args_to
  727. def before_hr(self, p):
  728. for script in self.ordered_scripts('before_hr'):
  729. try:
  730. script_args = p.script_args[script.args_from:script.args_to]
  731. script.before_hr(p, *script_args)
  732. except Exception:
  733. errors.report(f"Error running before_hr: {script.filename}", exc_info=True)
  734. def setup_scrips(self, p, *, is_ui=True):
  735. for script in self.ordered_scripts('setup'):
  736. if not is_ui and script.setup_for_ui_only:
  737. continue
  738. try:
  739. script_args = p.script_args[script.args_from:script.args_to]
  740. script.setup(p, *script_args)
  741. except Exception:
  742. errors.report(f"Error running setup: {script.filename}", exc_info=True)
  743. def set_named_arg(self, args, script_name, arg_elem_id, value, fuzzy=False):
  744. """Locate an arg of a specific script in script_args and set its value
  745. Args:
  746. args: all script args of process p, p.script_args
  747. script_name: the name target script name to
  748. arg_elem_id: the elem_id of the target arg
  749. value: the value to set
  750. fuzzy: if True, arg_elem_id can be a substring of the control.elem_id else exact match
  751. Returns:
  752. Updated script args
  753. when script_name in not found or arg_elem_id is not found in script controls, raise RuntimeError
  754. """
  755. script = next((x for x in self.scripts if x.name == script_name), None)
  756. if script is None:
  757. raise RuntimeError(f"script {script_name} not found")
  758. for i, control in enumerate(script.controls):
  759. if arg_elem_id in control.elem_id if fuzzy else arg_elem_id == control.elem_id:
  760. index = script.args_from + i
  761. if isinstance(args, tuple):
  762. return args[:index] + (value,) + args[index + 1:]
  763. elif isinstance(args, list):
  764. args[index] = value
  765. return args
  766. else:
  767. raise RuntimeError(f"args is not a list or tuple, but {type(args)}")
  768. raise RuntimeError(f"arg_elem_id {arg_elem_id} not found in script {script_name}")
  769. scripts_txt2img: ScriptRunner = None
  770. scripts_img2img: ScriptRunner = None
  771. scripts_postproc: scripts_postprocessing.ScriptPostprocessingRunner = None
  772. scripts_current: ScriptRunner = None
  773. def reload_script_body_only():
  774. cache = {}
  775. scripts_txt2img.reload_sources(cache)
  776. scripts_img2img.reload_sources(cache)
  777. reload_scripts = load_scripts # compatibility alias