Add job argument to State.begin()

This commit is contained in:
Aarni Koskela 2023-06-30 13:11:31 +03:00
parent b70001e618
commit f44feb6a10
6 changed files with 13 additions and 16 deletions

View File

@ -327,7 +327,7 @@ class Api:
p.outpath_grids = opts.outdir_txt2img_grids p.outpath_grids = opts.outdir_txt2img_grids
p.outpath_samples = opts.outdir_txt2img_samples p.outpath_samples = opts.outdir_txt2img_samples
shared.state.begin() shared.state.begin(job="scripts_txt2img")
if selectable_scripts is not None: if selectable_scripts is not None:
p.script_args = script_args p.script_args = script_args
processed = scripts.scripts_txt2img.run(p, *p.script_args) # Need to pass args as list here processed = scripts.scripts_txt2img.run(p, *p.script_args) # Need to pass args as list here
@ -384,7 +384,7 @@ class Api:
p.outpath_grids = opts.outdir_img2img_grids p.outpath_grids = opts.outdir_img2img_grids
p.outpath_samples = opts.outdir_img2img_samples p.outpath_samples = opts.outdir_img2img_samples
shared.state.begin() shared.state.begin(job="scripts_img2img")
if selectable_scripts is not None: if selectable_scripts is not None:
p.script_args = script_args p.script_args = script_args
processed = scripts.scripts_img2img.run(p, *p.script_args) # Need to pass args as list here processed = scripts.scripts_img2img.run(p, *p.script_args) # Need to pass args as list here
@ -599,7 +599,7 @@ class Api:
def create_embedding(self, args: dict): def create_embedding(self, args: dict):
try: try:
shared.state.begin() shared.state.begin(job="create_embedding")
filename = create_embedding(**args) # create empty embedding filename = create_embedding(**args) # create empty embedding
sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings() # reload embeddings so new one can be immediately used sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings() # reload embeddings so new one can be immediately used
shared.state.end() shared.state.end()
@ -610,7 +610,7 @@ class Api:
def create_hypernetwork(self, args: dict): def create_hypernetwork(self, args: dict):
try: try:
shared.state.begin() shared.state.begin(job="create_hypernetwork")
filename = create_hypernetwork(**args) # create empty embedding filename = create_hypernetwork(**args) # create empty embedding
shared.state.end() shared.state.end()
return models.CreateResponse(info=f"create hypernetwork filename: {filename}") return models.CreateResponse(info=f"create hypernetwork filename: {filename}")
@ -620,7 +620,7 @@ class Api:
def preprocess(self, args: dict): def preprocess(self, args: dict):
try: try:
shared.state.begin() shared.state.begin(job="preprocess")
preprocess(**args) # quick operation unless blip/booru interrogation is enabled preprocess(**args) # quick operation unless blip/booru interrogation is enabled
shared.state.end() shared.state.end()
return models.PreprocessResponse(info = 'preprocess complete') return models.PreprocessResponse(info = 'preprocess complete')
@ -636,7 +636,7 @@ class Api:
def train_embedding(self, args: dict): def train_embedding(self, args: dict):
try: try:
shared.state.begin() shared.state.begin(job="train_embedding")
apply_optimizations = shared.opts.training_xattention_optimizations apply_optimizations = shared.opts.training_xattention_optimizations
error = None error = None
filename = '' filename = ''
@ -657,7 +657,7 @@ class Api:
def train_hypernetwork(self, args: dict): def train_hypernetwork(self, args: dict):
try: try:
shared.state.begin() shared.state.begin(job="train_hypernetwork")
shared.loaded_hypernetworks = [] shared.loaded_hypernetworks = []
apply_optimizations = shared.opts.training_xattention_optimizations apply_optimizations = shared.opts.training_xattention_optimizations
error = None error = None

View File

@ -30,7 +30,7 @@ def wrap_gradio_gpu_call(func, extra_outputs=None):
id_task = None id_task = None
with queue_lock: with queue_lock:
shared.state.begin() shared.state.begin(job=id_task)
progress.start_task(id_task) progress.start_task(id_task)
try: try:

View File

@ -73,8 +73,7 @@ def to_half(tensor, enable):
def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae, discard_weights, save_metadata): def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae, discard_weights, save_metadata):
shared.state.begin() shared.state.begin(job="model-merge")
shared.state.job = 'model-merge'
def fail(message): def fail(message):
shared.state.textinfo = message shared.state.textinfo = message

View File

@ -184,8 +184,7 @@ class InterrogateModels:
def interrogate(self, pil_image): def interrogate(self, pil_image):
res = "" res = ""
shared.state.begin() shared.state.begin(job="interrogate")
shared.state.job = 'interrogate'
try: try:
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram: if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
lowvram.send_everything_to_cpu() lowvram.send_everything_to_cpu()

View File

@ -9,8 +9,7 @@ from modules.shared import opts
def run_postprocessing(extras_mode, image, image_folder, input_dir, output_dir, show_extras_results, *args, save_output: bool = True): def run_postprocessing(extras_mode, image, image_folder, input_dir, output_dir, show_extras_results, *args, save_output: bool = True):
devices.torch_gc() devices.torch_gc()
shared.state.begin() shared.state.begin(job="extras")
shared.state.job = 'extras'
image_data = [] image_data = []
image_names = [] image_names = []

View File

@ -173,7 +173,7 @@ class State:
return obj return obj
def begin(self): def begin(self, job: str = "(unknown)"):
self.sampling_step = 0 self.sampling_step = 0
self.job_count = -1 self.job_count = -1
self.processing_has_refined_job_count = False self.processing_has_refined_job_count = False
@ -187,7 +187,7 @@ class State:
self.interrupted = False self.interrupted = False
self.textinfo = None self.textinfo = None
self.time_start = time.time() self.time_start = time.time()
self.job = job
devices.torch_gc() devices.torch_gc()
def end(self): def end(self):