From b31f546516f013a15dadf4843e5ddb73e99e3d9e Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Sat, 15 Jul 2023 18:54:54 -0500 Subject: [PATCH] lint(api): use more accurate worker name throughout --- api/onnx_web/chain/base.py | 24 +++++------ api/onnx_web/chain/blend_img2img.py | 4 +- api/onnx_web/chain/blend_linear.py | 2 +- api/onnx_web/chain/blend_mask.py | 2 +- api/onnx_web/chain/correct_codeformer.py | 4 +- api/onnx_web/chain/persist_disk.py | 2 +- api/onnx_web/chain/persist_s3.py | 2 +- api/onnx_web/chain/reduce_crop.py | 2 +- api/onnx_web/chain/reduce_thumbnail.py | 2 +- api/onnx_web/chain/source_noise.py | 2 +- api/onnx_web/chain/source_s3.py | 2 +- api/onnx_web/chain/source_txt2img.py | 4 +- api/onnx_web/chain/source_url.py | 2 +- api/onnx_web/chain/stage.py | 2 +- api/onnx_web/chain/upscale_bsrgan.py | 4 +- api/onnx_web/chain/upscale_highres.py | 4 +- api/onnx_web/chain/upscale_outpaint.py | 4 +- api/onnx_web/chain/upscale_resrgan.py | 4 +- api/onnx_web/chain/upscale_simple.py | 2 +- .../chain/upscale_stable_diffusion.py | 4 +- api/onnx_web/chain/upscale_swinir.py | 4 +- api/onnx_web/diffusers/run.py | 40 +++++++++---------- api/onnx_web/transformers/run.py | 4 +- 23 files changed, 63 insertions(+), 63 deletions(-) diff --git a/api/onnx_web/chain/base.py b/api/onnx_web/chain/base.py index a3ffb54d..68405ff2 100644 --- a/api/onnx_web/chain/base.py +++ b/api/onnx_web/chain/base.py @@ -70,14 +70,14 @@ class ChainPipeline: def run( self, - job: WorkerContext, + worker: WorkerContext, server: ServerContext, params: ImageParams, sources: List[Image.Image], callback: Optional[ProgressCallback], **kwargs ) -> List[Image.Image]: - return self(job, server, params, sources=sources, callback=callback, **kwargs) + return self(worker, server, params, sources=sources, callback=callback, **kwargs) def stage(self, callback: BaseStage, params: StageParams, **kwargs): self.stages.append((callback, params, kwargs)) @@ -85,7 +85,7 @@ class ChainPipeline: def __call__( self, - job: WorkerContext, + worker: WorkerContext, server: ServerContext, params: ImageParams, sources: List[Image.Image], @@ -151,10 +151,10 @@ class ChainPipeline: tile_mask: Image.Image, dims: Tuple[int, int, int], ) -> Image.Image: - for i in range(job.retries): + for i in range(worker.retries): try: output_tile = stage_pipe.run( - job, + worker, server, stage_params, params, @@ -175,8 +175,8 @@ class ChainPipeline: i, ) server.cache.clear() - run_gc([job.get_device()]) - job.retries = job.retries - (i + 1) + run_gc([worker.get_device()]) + worker.retries = worker.retries - (i + 1) raise RetryException("exhausted retries on tile") @@ -193,10 +193,10 @@ class ChainPipeline: stage_sources = stage_outputs else: logger.debug("image within tile size of %s, running stage", tile) - for i in range(job.retries): + for i in range(worker.retries): try: stage_outputs = stage_pipe.run( - job, + worker, server, stage_params, params, @@ -213,10 +213,10 @@ class ChainPipeline: "error while running stage pipeline, retry %s of 3", i ) server.cache.clear() - run_gc([job.get_device()]) - job.retries = job.retries - (i + 1) + run_gc([worker.get_device()]) + worker.retries = worker.retries - (i + 1) - if job.retries <= 0: + if worker.retries <= 0: raise RetryException("exhausted retries on stage") logger.debug( diff --git a/api/onnx_web/chain/blend_img2img.py b/api/onnx_web/chain/blend_img2img.py index f4974c42..572ac3d3 100644 --- a/api/onnx_web/chain/blend_img2img.py +++ b/api/onnx_web/chain/blend_img2img.py @@ -20,7 +20,7 @@ class BlendImg2ImgStage(BaseStage): def run( self, - job: WorkerContext, + worker: WorkerContext, server: ServerContext, _stage: StageParams, params: ImageParams, @@ -45,7 +45,7 @@ class BlendImg2ImgStage(BaseStage): server, params, pipe_type, - job.get_device(), + worker.get_device(), inversions=inversions, loras=loras, ) diff --git a/api/onnx_web/chain/blend_linear.py b/api/onnx_web/chain/blend_linear.py index e56ba2e4..6317ef13 100644 --- a/api/onnx_web/chain/blend_linear.py +++ b/api/onnx_web/chain/blend_linear.py @@ -14,7 +14,7 @@ logger = getLogger(__name__) class BlendLinearStage(BaseStage): def run( self, - _job: WorkerContext, + _worker: WorkerContext, _server: ServerContext, _stage: StageParams, _params: ImageParams, diff --git a/api/onnx_web/chain/blend_mask.py b/api/onnx_web/chain/blend_mask.py index db956d42..1038d3ea 100644 --- a/api/onnx_web/chain/blend_mask.py +++ b/api/onnx_web/chain/blend_mask.py @@ -16,7 +16,7 @@ logger = getLogger(__name__) class BlendMaskStage(BaseStage): def run( self, - _job: WorkerContext, + _worker: WorkerContext, server: ServerContext, _stage: StageParams, _params: ImageParams, diff --git a/api/onnx_web/chain/correct_codeformer.py b/api/onnx_web/chain/correct_codeformer.py index e86629c1..121a5cb3 100644 --- a/api/onnx_web/chain/correct_codeformer.py +++ b/api/onnx_web/chain/correct_codeformer.py @@ -14,7 +14,7 @@ logger = getLogger(__name__) class CorrectCodeformerStage(BaseStage): def run( self, - job: WorkerContext, + worker: WorkerContext, _server: ServerContext, _stage: StageParams, _params: ImageParams, @@ -30,6 +30,6 @@ class CorrectCodeformerStage(BaseStage): upscale = upscale.with_args(**kwargs) - device = job.get_device() + device = worker.get_device() pipe = CodeFormer(upscale=upscale.face_outscale).to(device.torch_str()) return [pipe(source) for source in sources] diff --git a/api/onnx_web/chain/persist_disk.py b/api/onnx_web/chain/persist_disk.py index b5d24b96..794dbb37 100644 --- a/api/onnx_web/chain/persist_disk.py +++ b/api/onnx_web/chain/persist_disk.py @@ -15,7 +15,7 @@ logger = getLogger(__name__) class PersistDiskStage(BaseStage): def run( self, - _job: WorkerContext, + _worker: WorkerContext, server: ServerContext, _stage: StageParams, params: ImageParams, diff --git a/api/onnx_web/chain/persist_s3.py b/api/onnx_web/chain/persist_s3.py index a12e39df..27f4026f 100644 --- a/api/onnx_web/chain/persist_s3.py +++ b/api/onnx_web/chain/persist_s3.py @@ -16,7 +16,7 @@ logger = getLogger(__name__) class PersistS3Stage(BaseStage): def run( self, - _job: WorkerContext, + _worker: WorkerContext, server: ServerContext, _stage: StageParams, _params: ImageParams, diff --git a/api/onnx_web/chain/reduce_crop.py b/api/onnx_web/chain/reduce_crop.py index c784b1ea..2e258075 100644 --- a/api/onnx_web/chain/reduce_crop.py +++ b/api/onnx_web/chain/reduce_crop.py @@ -14,7 +14,7 @@ logger = getLogger(__name__) class ReduceCropStage(BaseStage): def run( self, - _job: WorkerContext, + _worker: WorkerContext, _server: ServerContext, _stage: StageParams, _params: ImageParams, diff --git a/api/onnx_web/chain/reduce_thumbnail.py b/api/onnx_web/chain/reduce_thumbnail.py index 3ae4f320..d7a0efee 100644 --- a/api/onnx_web/chain/reduce_thumbnail.py +++ b/api/onnx_web/chain/reduce_thumbnail.py @@ -14,7 +14,7 @@ logger = getLogger(__name__) class ReduceThumbnailStage(BaseStage): def run( self, - _job: WorkerContext, + _worker: WorkerContext, _server: ServerContext, _stage: StageParams, _params: ImageParams, diff --git a/api/onnx_web/chain/source_noise.py b/api/onnx_web/chain/source_noise.py index 27b12b03..ca878fb1 100644 --- a/api/onnx_web/chain/source_noise.py +++ b/api/onnx_web/chain/source_noise.py @@ -14,7 +14,7 @@ logger = getLogger(__name__) class SourceNoiseStage(BaseStage): def run( self, - _job: WorkerContext, + _worker: WorkerContext, _server: ServerContext, _stage: StageParams, _params: ImageParams, diff --git a/api/onnx_web/chain/source_s3.py b/api/onnx_web/chain/source_s3.py index c10a99d8..900270a3 100644 --- a/api/onnx_web/chain/source_s3.py +++ b/api/onnx_web/chain/source_s3.py @@ -16,7 +16,7 @@ logger = getLogger(__name__) class SourceS3Stage(BaseStage): def run( self, - _job: WorkerContext, + _worker: WorkerContext, _server: ServerContext, _stage: StageParams, _params: ImageParams, diff --git a/api/onnx_web/chain/source_txt2img.py b/api/onnx_web/chain/source_txt2img.py index ef470d9c..31a57815 100644 --- a/api/onnx_web/chain/source_txt2img.py +++ b/api/onnx_web/chain/source_txt2img.py @@ -25,7 +25,7 @@ class SourceTxt2ImgStage(BaseStage): def run( self, - job: WorkerContext, + worker: WorkerContext, server: ServerContext, _stage: StageParams, params: ImageParams, @@ -67,7 +67,7 @@ class SourceTxt2ImgStage(BaseStage): server, params, pipe_type, - job.get_device(), + worker.get_device(), inversions=inversions, loras=loras, ) diff --git a/api/onnx_web/chain/source_url.py b/api/onnx_web/chain/source_url.py index 86c4f67e..3883ef74 100644 --- a/api/onnx_web/chain/source_url.py +++ b/api/onnx_web/chain/source_url.py @@ -16,7 +16,7 @@ logger = getLogger(__name__) class SourceURLStage(BaseStage): def run( self, - _job: WorkerContext, + _worker: WorkerContext, _server: ServerContext, _stage: StageParams, _params: ImageParams, diff --git a/api/onnx_web/chain/stage.py b/api/onnx_web/chain/stage.py index 11c61774..781b65de 100644 --- a/api/onnx_web/chain/stage.py +++ b/api/onnx_web/chain/stage.py @@ -12,7 +12,7 @@ class BaseStage: def run( self, - job: WorkerContext, + worker: WorkerContext, server: ServerContext, stage: StageParams, _params: ImageParams, diff --git a/api/onnx_web/chain/upscale_bsrgan.py b/api/onnx_web/chain/upscale_bsrgan.py index 852d5806..1f678237 100644 --- a/api/onnx_web/chain/upscale_bsrgan.py +++ b/api/onnx_web/chain/upscale_bsrgan.py @@ -50,7 +50,7 @@ class UpscaleBSRGANStage(BaseStage): def run( self, - job: WorkerContext, + worker: WorkerContext, server: ServerContext, stage: StageParams, _params: ImageParams, @@ -67,7 +67,7 @@ class UpscaleBSRGANStage(BaseStage): return sources logger.info("upscaling with BSRGAN model: %s", upscale.upscale_model) - device = job.get_device() + device = worker.get_device() bsrgan = self.load(server, stage, upscale, device) outputs = [] diff --git a/api/onnx_web/chain/upscale_highres.py b/api/onnx_web/chain/upscale_highres.py index 4ae96014..e19f75fb 100644 --- a/api/onnx_web/chain/upscale_highres.py +++ b/api/onnx_web/chain/upscale_highres.py @@ -16,7 +16,7 @@ logger = getLogger(__name__) class UpscaleHighresStage(BaseStage): def run( self, - job: WorkerContext, + worker: WorkerContext, server: ServerContext, stage: StageParams, params: ImageParams, @@ -35,7 +35,7 @@ class UpscaleHighresStage(BaseStage): return [ chain( - job, + worker, server, params, source, diff --git a/api/onnx_web/chain/upscale_outpaint.py b/api/onnx_web/chain/upscale_outpaint.py index ec1a9f1c..21b8e544 100644 --- a/api/onnx_web/chain/upscale_outpaint.py +++ b/api/onnx_web/chain/upscale_outpaint.py @@ -28,7 +28,7 @@ class UpscaleOutpaintStage(BaseStage): def run( self, - job: WorkerContext, + worker: WorkerContext, server: ServerContext, stage: StageParams, params: ImageParams, @@ -55,7 +55,7 @@ class UpscaleOutpaintStage(BaseStage): server, params, pipe_type, - job.get_device(), + worker.get_device(), inversions=inversions, loras=loras, ) diff --git a/api/onnx_web/chain/upscale_resrgan.py b/api/onnx_web/chain/upscale_resrgan.py index a75cc609..e680af53 100644 --- a/api/onnx_web/chain/upscale_resrgan.py +++ b/api/onnx_web/chain/upscale_resrgan.py @@ -73,7 +73,7 @@ class UpscaleRealESRGANStage(BaseStage): def run( self, - job: WorkerContext, + worker: WorkerContext, server: ServerContext, stage: StageParams, _params: ImageParams, @@ -89,7 +89,7 @@ class UpscaleRealESRGANStage(BaseStage): for source in sources: output = np.array(source) upsampler = self.load( - server, upscale, job.get_device(), tile=stage.tile_size + server, upscale, worker.get_device(), tile=stage.tile_size ) output, _ = upsampler.enhance(output, outscale=upscale.outscale) diff --git a/api/onnx_web/chain/upscale_simple.py b/api/onnx_web/chain/upscale_simple.py index 5e5a7e75..7dd44200 100644 --- a/api/onnx_web/chain/upscale_simple.py +++ b/api/onnx_web/chain/upscale_simple.py @@ -14,7 +14,7 @@ logger = getLogger(__name__) class UpscaleSimpleStage(BaseStage): def run( self, - _job: WorkerContext, + _worker: WorkerContext, _server: ServerContext, _stage: StageParams, _params: ImageParams, diff --git a/api/onnx_web/chain/upscale_stable_diffusion.py b/api/onnx_web/chain/upscale_stable_diffusion.py index b2fd78dc..9d5a7b32 100644 --- a/api/onnx_web/chain/upscale_stable_diffusion.py +++ b/api/onnx_web/chain/upscale_stable_diffusion.py @@ -18,7 +18,7 @@ logger = getLogger(__name__) class UpscaleStableDiffusionStage(BaseStage): def run( self, - job: WorkerContext, + worker: WorkerContext, server: ServerContext, _stage: StageParams, params: ImageParams, @@ -43,7 +43,7 @@ class UpscaleStableDiffusionStage(BaseStage): server, params, "upscale", - job.get_device(), + worker.get_device(), model=path.join(server.model_path, upscale.upscale_model), ) generator = torch.manual_seed(params.seed) diff --git a/api/onnx_web/chain/upscale_swinir.py b/api/onnx_web/chain/upscale_swinir.py index bba808d1..f78ee344 100644 --- a/api/onnx_web/chain/upscale_swinir.py +++ b/api/onnx_web/chain/upscale_swinir.py @@ -50,7 +50,7 @@ class UpscaleSwinIRStage(BaseStage): def run( self, - job: WorkerContext, + worker: WorkerContext, server: ServerContext, stage: StageParams, _params: ImageParams, @@ -67,7 +67,7 @@ class UpscaleSwinIRStage(BaseStage): return sources logger.info("correcting faces with SwinIR model: %s", upscale.upscale_model) - device = job.get_device() + device = worker.get_device() swinir = self.load(server, stage, upscale, device) outputs = [] diff --git a/api/onnx_web/diffusers/run.py b/api/onnx_web/diffusers/run.py index 36df4cd5..6e2e7497 100644 --- a/api/onnx_web/diffusers/run.py +++ b/api/onnx_web/diffusers/run.py @@ -34,7 +34,7 @@ logger = getLogger(__name__) def run_txt2img_pipeline( - job: WorkerContext, + worker: WorkerContext, server: ServerContext, params: ImageParams, size: Size, @@ -83,8 +83,8 @@ def run_txt2img_pipeline( # run and save latents = get_latents_from_seed(params.seed, size, batch=params.batch) - progress = job.get_progress_callback() - images = chain.run(job, server, params, [], callback=progress, latents=latents) + progress = worker.get_progress_callback() + images = chain.run(worker, server, params, [], callback=progress, latents=latents) _pairs, loras, inversions, _rest = parse_prompt(params) @@ -102,7 +102,7 @@ def run_txt2img_pipeline( ) # clean up - run_gc([job.get_device()]) + run_gc([worker.get_device()]) # notify the user show_system_toast(f"finished txt2img job: {dest}") @@ -110,7 +110,7 @@ def run_txt2img_pipeline( def run_img2img_pipeline( - job: WorkerContext, + worker: WorkerContext, server: ServerContext, params: ImageParams, outputs: List[str], @@ -175,8 +175,8 @@ def run_img2img_pipeline( ) # run and append the filtered source - progress = job.get_progress_callback() - images = chain(job, server, params, [source], callback=progress) + progress = worker.get_progress_callback() + images = chain(worker, server, params, [source], callback=progress) if source_filter is not None and source_filter != "none": images.append(source) @@ -199,7 +199,7 @@ def run_img2img_pipeline( ) # clean up - run_gc([job.get_device()]) + run_gc([worker.get_device()]) # notify the user show_system_toast(f"finished img2img job: {dest}") @@ -207,7 +207,7 @@ def run_img2img_pipeline( def run_inpaint_pipeline( - job: WorkerContext, + worker: WorkerContext, server: ServerContext, params: ImageParams, size: Size, @@ -353,8 +353,8 @@ def run_inpaint_pipeline( # run and save latents = get_latents_from_seed(params.seed, size, batch=params.batch) - progress = job.get_progress_callback() - images = chain(job, server, params, [source], callback=progress, latents=latents) + progress = worker.get_progress_callback() + images = chain(worker, server, params, [source], callback=progress, latents=latents) _pairs, loras, inversions, _rest = parse_prompt(params) for image, output in zip(images, outputs): @@ -378,7 +378,7 @@ def run_inpaint_pipeline( # clean up del image - run_gc([job.get_device()]) + run_gc([worker.get_device()]) # notify the user show_system_toast(f"finished inpaint job: {dest}") @@ -386,7 +386,7 @@ def run_inpaint_pipeline( def run_upscale_pipeline( - job: WorkerContext, + worker: WorkerContext, server: ServerContext, params: ImageParams, size: Size, @@ -427,8 +427,8 @@ def run_upscale_pipeline( ) # run and save - progress = job.get_progress_callback() - images = chain(job, server, params, [source], callback=progress) + progress = worker.get_progress_callback() + images = chain(worker, server, params, [source], callback=progress) _pairs, loras, inversions, _rest = parse_prompt(params) for image, output in zip(images, outputs): @@ -445,7 +445,7 @@ def run_upscale_pipeline( # clean up del image - run_gc([job.get_device()]) + run_gc([worker.get_device()]) # notify the user show_system_toast(f"finished upscale job: {dest}") @@ -453,7 +453,7 @@ def run_upscale_pipeline( def run_blend_pipeline( - job: WorkerContext, + worker: WorkerContext, server: ServerContext, params: ImageParams, size: Size, @@ -477,15 +477,15 @@ def run_blend_pipeline( ) # run and save - progress = job.get_progress_callback() - images = chain(job, server, params, sources, callback=progress) + progress = worker.get_progress_callback() + images = chain(worker, server, params, sources, callback=progress) for image, output in zip(images, outputs): dest = save_image(server, output, image, params, size, upscale=upscale) # clean up del image - run_gc([job.get_device()]) + run_gc([worker.get_device()]) # notify the user show_system_toast(f"finished blend job: {dest}") diff --git a/api/onnx_web/transformers/run.py b/api/onnx_web/transformers/run.py index 25de1071..eb789639 100644 --- a/api/onnx_web/transformers/run.py +++ b/api/onnx_web/transformers/run.py @@ -8,7 +8,7 @@ logger = getLogger(__name__) def run_txt2txt_pipeline( - job: WorkerContext, + worker: WorkerContext, _server: ServerContext, params: ImageParams, _size: Size, @@ -20,7 +20,7 @@ def run_txt2txt_pipeline( model = "EleutherAI/gpt-j-6B" tokens = 1024 - device = job.get_device() + device = worker.get_device() pipe = GPTJForCausalLM.from_pretrained(model).to(device.torch_str()) tokenizer = AutoTokenizer.from_pretrained(model)