From fb58fa62400ba0e2cbe8703796d5573f5f91c398 Mon Sep 17 00:00:00 2001 From: EllangoK Date: Sat, 28 Jan 2023 15:37:01 -0500 Subject: [PATCH 01/10] xyz plot now saves sub grids if opts.grid_save also fixed no draw legend for z grid --- scripts/xyz_grid.py | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py index 3df40483..3122f6f6 100644 --- a/scripts/xyz_grid.py +++ b/scripts/xyz_grid.py @@ -286,23 +286,24 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend print("Unexpected error: draw_xyz_grid failed to return even a single processed image") return Processed(p, []) - grids = [None] * len(zs) + sub_grids = [None] * len(zs) for i in range(len(zs)): start_index = i * len(xs) * len(ys) end_index = start_index + len(xs) * len(ys) grid = images.image_grid(image_cache[start_index:end_index], rows=len(ys)) if draw_legend: grid = images.draw_grid_annotations(grid, cell_size[0], cell_size[1], hor_texts, ver_texts) - - grids[i] = grid + sub_grids[i] = grid if include_sub_grids and len(zs) > 1: processed_result.images.insert(i+1, grid) - original_grid_size = grids[0].size - grids = images.image_grid(grids, rows=1) - processed_result.images[0] = images.draw_grid_annotations(grids, original_grid_size[0], original_grid_size[1], title_texts, [[images.GridAnnotation()]]) + sub_grid_size = sub_grids[0].size + z_grid = images.image_grid(sub_grids, rows=1) + if draw_legend: + z_grid = images.draw_grid_annotations(z_grid, sub_grid_size[0], sub_grid_size[1], title_texts, [[images.GridAnnotation()]]) + processed_result.images[0] = z_grid - return processed_result + return processed_result, sub_grids class SharedSettingsStackHelper(object): @@ -576,7 +577,7 @@ class Script(scripts.Script): return res with SharedSettingsStackHelper(): - processed = draw_xyz_grid( + processed, sub_grids = draw_xyz_grid( p, xs=xs, ys=ys, @@ -592,6 +593,10 @@ class Script(scripts.Script): second_axes_processed=second_axes_processed ) + if opts.grid_save and len(sub_grids) > 1: + for sub_grid in sub_grids: + images.save_image(sub_grid, p.outpath_grids, "xyz_grid", info=grid_infotext[0], extension=opts.grid_format, prompt=p.prompt, seed=processed.seed, grid=True, p=p) + if opts.grid_save: images.save_image(processed.images[0], p.outpath_grids, "xyz_grid", info=grid_infotext[0], extension=opts.grid_format, prompt=p.prompt, seed=processed.seed, grid=True, p=p) From 920fe8057cb325e9835f70c0389499c51cbdd3b5 Mon Sep 17 00:00:00 2001 From: EllangoK Date: Sun, 29 Jan 2023 03:36:16 -0500 Subject: [PATCH 02/10] fixes #7284 btn unbound error --- modules/ui.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modules/ui.py b/modules/ui.py index f1195692..7e193240 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -466,8 +466,8 @@ def create_ui(): width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512, elem_id="txt2img_width") height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512, elem_id="txt2img_height") + res_switch_btn = ToolButton(value=switch_values_symbol, elem_id="txt2img_res_switch_btn") if opts.dimensions_and_batch_together: - res_switch_btn = ToolButton(value=switch_values_symbol, elem_id="txt2img_res_switch_btn") with gr.Column(elem_id="txt2img_column_batch"): batch_count = gr.Slider(minimum=1, step=1, label='Batch count', value=1, elem_id="txt2img_batch_count") batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1, elem_id="txt2img_batch_size") @@ -737,8 +737,8 @@ def create_ui(): width = gr.Slider(minimum=64, maximum=2048, step=8, label="Width", value=512, elem_id="img2img_width") height = gr.Slider(minimum=64, maximum=2048, step=8, label="Height", value=512, elem_id="img2img_height") + res_switch_btn = ToolButton(value=switch_values_symbol, elem_id="img2img_res_switch_btn") if opts.dimensions_and_batch_together: - res_switch_btn = ToolButton(value=switch_values_symbol, elem_id="img2img_res_switch_btn") with gr.Column(elem_id="img2img_column_batch"): batch_count = gr.Slider(minimum=1, step=1, label='Batch count', value=1, elem_id="img2img_batch_count") batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1, elem_id="img2img_batch_size") From 7c53f81caf817a7e7dc9c2fafebfcce269ecb1d7 Mon Sep 17 00:00:00 2001 From: Yevhenii Hurin Date: Sun, 29 Jan 2023 15:29:03 +0200 Subject: [PATCH 03/10] Prompt selector for Prompt Matrix script --- scripts/prompt_matrix.py | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/scripts/prompt_matrix.py b/scripts/prompt_matrix.py index dd95e588..702870ce 100644 --- a/scripts/prompt_matrix.py +++ b/scripts/prompt_matrix.py @@ -10,7 +10,7 @@ from modules import images from modules.processing import process_images, Processed from modules.shared import opts, cmd_opts, state import modules.sd_samplers - +from pprint import pprint def draw_xy_grid(xs, ys, x_label, y_label, cell): res = [] @@ -44,16 +44,23 @@ class Script(scripts.Script): def title(self): return "Prompt matrix" - def ui(self, is_img2img): + def ui(self, is_img2img): put_at_start = gr.Checkbox(label='Put variable parts at start of prompt', value=False, elem_id=self.elem_id("put_at_start")) different_seeds = gr.Checkbox(label='Use different seed for each picture', value=False, elem_id=self.elem_id("different_seeds")) + # Radio buttons for selecting the prompt between positive and negative + prompt_type = gr.Radio(["positive", "negative"], label="Select prompt", elem_id=self.elem_id("prompt_type"), value="positive") - return [put_at_start, different_seeds] + return [put_at_start, different_seeds, prompt_type] - def run(self, p, put_at_start, different_seeds): + def run(self, p, put_at_start, different_seeds, prompt_type): modules.processing.fix_seed(p) + # Raise error if promp type is not positive or negative + if prompt_type not in ["positive", "negative"]: + raise ValueError(f"Unknown prompt type {prompt_type}") - original_prompt = p.prompt[0] if type(p.prompt) == list else p.prompt + prompt = p.prompt if prompt_type == "positive" else p.negative_prompt + original_prompt = prompt[0] if type(prompt) == list else prompt + positive_prompt = p.prompt[0] if type(p.prompt) == list else p.prompt all_prompts = [] prompt_matrix_parts = original_prompt.split("|") @@ -73,9 +80,12 @@ class Script(scripts.Script): print(f"Prompt matrix will create {len(all_prompts)} images using a total of {p.n_iter} batches.") - p.prompt = all_prompts + if prompt_type == "positive": + p.prompt = all_prompts + else: + p.negative_prompt = all_prompts p.seed = [p.seed + (i if different_seeds else 0) for i in range(len(all_prompts))] - p.prompt_for_display = original_prompt + p.prompt_for_display = positive_prompt processed = process_images(p) grid = images.image_grid(processed.images, p.batch_size, rows=1 << ((len(prompt_matrix_parts) - 1) // 2)) From edabd927296d389856ffc87c7789111a21a051a2 Mon Sep 17 00:00:00 2001 From: Yevhenii Hurin Date: Sun, 29 Jan 2023 16:05:59 +0200 Subject: [PATCH 04/10] Add delimiter selector to the Prompt Matrix script --- scripts/prompt_matrix.py | 47 ++++++++++++++++++++++++++++------------ 1 file changed, 33 insertions(+), 14 deletions(-) diff --git a/scripts/prompt_matrix.py b/scripts/prompt_matrix.py index 702870ce..89db9e63 100644 --- a/scripts/prompt_matrix.py +++ b/scripts/prompt_matrix.py @@ -12,6 +12,7 @@ from modules.shared import opts, cmd_opts, state import modules.sd_samplers from pprint import pprint + def draw_xy_grid(xs, ys, x_label, y_label, cell): res = [] @@ -33,7 +34,8 @@ def draw_xy_grid(xs, ys, x_label, y_label, cell): res.append(processed.images[0]) grid = images.image_grid(res, rows=len(ys)) - grid = images.draw_grid_annotations(grid, res[0].width, res[0].height, hor_texts, ver_texts) + grid = images.draw_grid_annotations( + grid, res[0].width, res[0].height, hor_texts, ver_texts) first_processed.images = [grid] @@ -45,56 +47,73 @@ class Script(scripts.Script): return "Prompt matrix" def ui(self, is_img2img): - put_at_start = gr.Checkbox(label='Put variable parts at start of prompt', value=False, elem_id=self.elem_id("put_at_start")) - different_seeds = gr.Checkbox(label='Use different seed for each picture', value=False, elem_id=self.elem_id("different_seeds")) + put_at_start = gr.Checkbox(label='Put variable parts at start of prompt', + value=False, elem_id=self.elem_id("put_at_start")) + different_seeds = gr.Checkbox( + label='Use different seed for each picture', value=False, elem_id=self.elem_id("different_seeds")) # Radio buttons for selecting the prompt between positive and negative - prompt_type = gr.Radio(["positive", "negative"], label="Select prompt", elem_id=self.elem_id("prompt_type"), value="positive") + prompt_type = gr.Radio(["positive", "negative"], label="Select prompt", + elem_id=self.elem_id("prompt_type"), value="positive") + # Radio buttons for selecting the delimiter to use in the resulting prompt + variations_delimiter = gr.Radio(["comma", "space"], label="Select delimiter", elem_id=self.elem_id( + "variations_delimiter"), value="comma") + return [put_at_start, different_seeds, prompt_type, variations_delimiter] - return [put_at_start, different_seeds, prompt_type] - - def run(self, p, put_at_start, different_seeds, prompt_type): + def run(self, p, put_at_start, different_seeds, prompt_type, variations_delimiter): modules.processing.fix_seed(p) # Raise error if promp type is not positive or negative if prompt_type not in ["positive", "negative"]: raise ValueError(f"Unknown prompt type {prompt_type}") + # Raise error if variations delimiter is not comma or space + if variations_delimiter not in ["comma", "space"]: + raise ValueError( + f"Unknown variations delimiter {variations_delimiter}") prompt = p.prompt if prompt_type == "positive" else p.negative_prompt original_prompt = prompt[0] if type(prompt) == list else prompt positive_prompt = p.prompt[0] if type(p.prompt) == list else p.prompt + delimiter = ", " if variations_delimiter == "comma" else " " + all_prompts = [] prompt_matrix_parts = original_prompt.split("|") combination_count = 2 ** (len(prompt_matrix_parts) - 1) for combination_num in range(combination_count): - selected_prompts = [text.strip().strip(',') for n, text in enumerate(prompt_matrix_parts[1:]) if combination_num & (1 << n)] + selected_prompts = [text.strip().strip(',') for n, text in enumerate( + prompt_matrix_parts[1:]) if combination_num & (1 << n)] if put_at_start: selected_prompts = selected_prompts + [prompt_matrix_parts[0]] else: selected_prompts = [prompt_matrix_parts[0]] + selected_prompts - all_prompts.append(", ".join(selected_prompts)) + all_prompts.append(delimiter.join(selected_prompts)) p.n_iter = math.ceil(len(all_prompts) / p.batch_size) p.do_not_save_grid = True - print(f"Prompt matrix will create {len(all_prompts)} images using a total of {p.n_iter} batches.") + print( + f"Prompt matrix will create {len(all_prompts)} images using a total of {p.n_iter} batches.") if prompt_type == "positive": p.prompt = all_prompts else: p.negative_prompt = all_prompts - p.seed = [p.seed + (i if different_seeds else 0) for i in range(len(all_prompts))] + p.seed = [p.seed + (i if different_seeds else 0) + for i in range(len(all_prompts))] p.prompt_for_display = positive_prompt processed = process_images(p) - grid = images.image_grid(processed.images, p.batch_size, rows=1 << ((len(prompt_matrix_parts) - 1) // 2)) - grid = images.draw_prompt_matrix(grid, p.width, p.height, prompt_matrix_parts) + grid = images.image_grid(processed.images, p.batch_size, rows=1 << ( + (len(prompt_matrix_parts) - 1) // 2)) + grid = images.draw_prompt_matrix( + grid, p.width, p.height, prompt_matrix_parts) processed.images.insert(0, grid) processed.index_of_first_image = 1 processed.infotexts.insert(0, processed.infotexts[0]) if opts.grid_save: - images.save_image(processed.images[0], p.outpath_grids, "prompt_matrix", extension=opts.grid_format, prompt=original_prompt, seed=processed.seed, grid=True, p=p) + images.save_image(processed.images[0], p.outpath_grids, "prompt_matrix", + extension=opts.grid_format, prompt=original_prompt, seed=processed.seed, grid=True, p=p) return processed From 5997457fd48c9c0fc31c9d96bf0b9c217585c526 Mon Sep 17 00:00:00 2001 From: Yevhenii Hurin Date: Sun, 29 Jan 2023 16:23:29 +0200 Subject: [PATCH 05/10] Compact options UI for Prompt Matrix --- scripts/prompt_matrix.py | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/scripts/prompt_matrix.py b/scripts/prompt_matrix.py index 89db9e63..03212b31 100644 --- a/scripts/prompt_matrix.py +++ b/scripts/prompt_matrix.py @@ -47,16 +47,23 @@ class Script(scripts.Script): return "Prompt matrix" def ui(self, is_img2img): - put_at_start = gr.Checkbox(label='Put variable parts at start of prompt', - value=False, elem_id=self.elem_id("put_at_start")) - different_seeds = gr.Checkbox( - label='Use different seed for each picture', value=False, elem_id=self.elem_id("different_seeds")) - # Radio buttons for selecting the prompt between positive and negative - prompt_type = gr.Radio(["positive", "negative"], label="Select prompt", - elem_id=self.elem_id("prompt_type"), value="positive") - # Radio buttons for selecting the delimiter to use in the resulting prompt - variations_delimiter = gr.Radio(["comma", "space"], label="Select delimiter", elem_id=self.elem_id( - "variations_delimiter"), value="comma") + gr.HTML('
') + with gr.Row(): + with gr.Column(): + put_at_start = gr.Checkbox(label='Put variable parts at start of prompt', + value=False, elem_id=self.elem_id("put_at_start")) + with gr.Column(): + # Radio buttons for selecting the prompt between positive and negative + prompt_type = gr.Radio(["positive", "negative"], label="Select prompt", + elem_id=self.elem_id("prompt_type"), value="positive") + with gr.Row(): + with gr.Column(): + different_seeds = gr.Checkbox( + label='Use different seed for each picture', value=False, elem_id=self.elem_id("different_seeds")) + with gr.Column(): + # Radio buttons for selecting the delimiter to use in the resulting prompt + variations_delimiter = gr.Radio(["comma", "space"], label="Select delimiter", elem_id=self.elem_id( + "variations_delimiter"), value="comma") return [put_at_start, different_seeds, prompt_type, variations_delimiter] def run(self, p, put_at_start, different_seeds, prompt_type, variations_delimiter): From 1e2b10d2dcdf41a6cce0c525c85ebd42a521e0f1 Mon Sep 17 00:00:00 2001 From: Yevhenii Hurin Date: Sun, 29 Jan 2023 17:14:46 +0200 Subject: [PATCH 06/10] Cleanup changes made by formatter --- scripts/prompt_matrix.py | 25 ++++++++----------------- 1 file changed, 8 insertions(+), 17 deletions(-) diff --git a/scripts/prompt_matrix.py b/scripts/prompt_matrix.py index 03212b31..de921ea8 100644 --- a/scripts/prompt_matrix.py +++ b/scripts/prompt_matrix.py @@ -10,7 +10,6 @@ from modules import images from modules.processing import process_images, Processed from modules.shared import opts, cmd_opts, state import modules.sd_samplers -from pprint import pprint def draw_xy_grid(xs, ys, x_label, y_label, cell): @@ -34,8 +33,7 @@ def draw_xy_grid(xs, ys, x_label, y_label, cell): res.append(processed.images[0]) grid = images.image_grid(res, rows=len(ys)) - grid = images.draw_grid_annotations( - grid, res[0].width, res[0].height, hor_texts, ver_texts) + grid = images.draw_grid_annotations(grid, res[0].width, res[0].height, hor_texts, ver_texts) first_processed.images = [grid] @@ -73,8 +71,7 @@ class Script(scripts.Script): raise ValueError(f"Unknown prompt type {prompt_type}") # Raise error if variations delimiter is not comma or space if variations_delimiter not in ["comma", "space"]: - raise ValueError( - f"Unknown variations delimiter {variations_delimiter}") + raise ValueError(f"Unknown variations delimiter {variations_delimiter}") prompt = p.prompt if prompt_type == "positive" else p.negative_prompt original_prompt = prompt[0] if type(prompt) == list else prompt @@ -86,8 +83,7 @@ class Script(scripts.Script): prompt_matrix_parts = original_prompt.split("|") combination_count = 2 ** (len(prompt_matrix_parts) - 1) for combination_num in range(combination_count): - selected_prompts = [text.strip().strip(',') for n, text in enumerate( - prompt_matrix_parts[1:]) if combination_num & (1 << n)] + selected_prompts = [text.strip().strip(',') for n, text in enumerate(prompt_matrix_parts[1:]) if combination_num & (1 << n)] if put_at_start: selected_prompts = selected_prompts + [prompt_matrix_parts[0]] @@ -99,28 +95,23 @@ class Script(scripts.Script): p.n_iter = math.ceil(len(all_prompts) / p.batch_size) p.do_not_save_grid = True - print( - f"Prompt matrix will create {len(all_prompts)} images using a total of {p.n_iter} batches.") + print(f"Prompt matrix will create {len(all_prompts)} images using a total of {p.n_iter} batches.") if prompt_type == "positive": p.prompt = all_prompts else: p.negative_prompt = all_prompts - p.seed = [p.seed + (i if different_seeds else 0) - for i in range(len(all_prompts))] + p.seed = [p.seed + (i if different_seeds else 0) for i in range(len(all_prompts))] p.prompt_for_display = positive_prompt processed = process_images(p) - grid = images.image_grid(processed.images, p.batch_size, rows=1 << ( - (len(prompt_matrix_parts) - 1) // 2)) - grid = images.draw_prompt_matrix( - grid, p.width, p.height, prompt_matrix_parts) + grid = images.image_grid(processed.images, p.batch_size, rows=1 << ((len(prompt_matrix_parts) - 1) // 2)) + grid = images.draw_prompt_matrix(grid, p.width, p.height, prompt_matrix_parts) processed.images.insert(0, grid) processed.index_of_first_image = 1 processed.infotexts.insert(0, processed.infotexts[0]) if opts.grid_save: - images.save_image(processed.images[0], p.outpath_grids, "prompt_matrix", - extension=opts.grid_format, prompt=original_prompt, seed=processed.seed, grid=True, p=p) + images.save_image(processed.images[0], p.outpath_grids, "prompt_matrix", extension=opts.grid_format, prompt=original_prompt, seed=processed.seed, grid=True, p=p) return processed From bfe7e7f15fbceccd016957769cd5b5a26c82a45b Mon Sep 17 00:00:00 2001 From: Piotr Pomierski Date: Tue, 31 Jan 2023 01:51:07 +0100 Subject: [PATCH 07/10] Fix missing tooltip for 'Clear prompt' button --- javascript/hints.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/javascript/hints.js b/javascript/hints.js index 7b60b25e..75792d0d 100644 --- a/javascript/hints.js +++ b/javascript/hints.js @@ -17,7 +17,7 @@ titles = { "\u2199\ufe0f": "Read generation parameters from prompt or last generation if prompt is empty into user interface.", "\u{1f4c2}": "Open images output directory", "\u{1f4be}": "Save style", - "\U0001F5D1": "Clear prompt", + "\u{1f5d1}": "Clear prompt", "\u{1f4cb}": "Apply selected styles to current prompt", "\u{1f4d2}": "Paste available values into the field", "\u{1f3b4}": "Show extra networks", From 0426b3478937e54446337cf435ed3f548688b120 Mon Sep 17 00:00:00 2001 From: Joey Sanchez Date: Mon, 30 Jan 2023 21:46:13 -0500 Subject: [PATCH 08/10] Adding default true to use_original_name_batch as images should by default hold the same name to help keep sequenced images in their correct order --- modules/shared.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/shared.py b/modules/shared.py index 69634fd8..5600d480 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -327,7 +327,7 @@ options_templates.update(options_section(('saving-images', "Saving images/grids" "jpeg_quality": OptionInfo(80, "Quality for saved jpeg images", gr.Slider, {"minimum": 1, "maximum": 100, "step": 1}), "export_for_4chan": OptionInfo(True, "If PNG image is larger than 4MB or any dimension is larger than 4000, downscale and save copy as JPG"), - "use_original_name_batch": OptionInfo(False, "Use original name for output filename during batch process in extras tab"), + "use_original_name_batch": OptionInfo(True, "Use original name for output filename during batch process in extras tab"), "use_upscaler_name_as_suffix": OptionInfo(False, "Use upscaler name as filename suffix in the extras tab"), "save_selected_only": OptionInfo(True, "When using 'Save' button, only save a single selected image"), "do_not_add_watermark": OptionInfo(False, "Do not add watermark to images"), From 7738c057ce938ca5c5a53a95e2023d3bcf14f06a Mon Sep 17 00:00:00 2001 From: brkirch Date: Wed, 1 Feb 2023 05:23:58 -0500 Subject: [PATCH 09/10] MPS fix is still needed :( Apparently I did not test with large enough images to trigger the bug with torch.narrow on MPS --- modules/devices.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/modules/devices.py b/modules/devices.py index 655ca1d3..f4afb897 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -207,3 +207,6 @@ if has_mps(): cumsum_needs_bool_fix = not torch.BoolTensor([True,True]).to(device=torch.device("mps"), dtype=torch.int64).equal(torch.BoolTensor([True,False]).to(torch.device("mps")).cumsum(0)) torch.cumsum = lambda input, *args, **kwargs: ( cumsum_fix(input, orig_cumsum, *args, **kwargs) ) torch.Tensor.cumsum = lambda self, *args, **kwargs: ( cumsum_fix(self, orig_Tensor_cumsum, *args, **kwargs) ) + orig_narrow = torch.narrow + torch.narrow = lambda *args, **kwargs: ( orig_narrow(*args, **kwargs).clone() ) + From 2217331cd1245d0bdda786a5dcaf4f7b843bd7e4 Mon Sep 17 00:00:00 2001 From: brkirch Date: Wed, 1 Feb 2023 06:20:19 -0500 Subject: [PATCH 10/10] Refactor MPS fixes to CondFunc --- modules/devices.py | 50 +++++++++++++--------------------------------- 1 file changed, 14 insertions(+), 36 deletions(-) diff --git a/modules/devices.py b/modules/devices.py index f4afb897..919048d0 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -2,6 +2,7 @@ import sys, os, shlex import contextlib import torch from modules import errors +from modules.sd_hijack_utils import CondFunc from packaging import version @@ -156,36 +157,7 @@ def test_for_nans(x, where): raise NansException(message) -# MPS workaround for https://github.com/pytorch/pytorch/issues/79383 -orig_tensor_to = torch.Tensor.to -def tensor_to_fix(self, *args, **kwargs): - if self.device.type != 'mps' and \ - ((len(args) > 0 and isinstance(args[0], torch.device) and args[0].type == 'mps') or \ - (isinstance(kwargs.get('device'), torch.device) and kwargs['device'].type == 'mps')): - self = self.contiguous() - return orig_tensor_to(self, *args, **kwargs) - - -# MPS workaround for https://github.com/pytorch/pytorch/issues/80800 -orig_layer_norm = torch.nn.functional.layer_norm -def layer_norm_fix(*args, **kwargs): - if len(args) > 0 and isinstance(args[0], torch.Tensor) and args[0].device.type == 'mps': - args = list(args) - args[0] = args[0].contiguous() - return orig_layer_norm(*args, **kwargs) - - -# MPS workaround for https://github.com/pytorch/pytorch/issues/90532 -orig_tensor_numpy = torch.Tensor.numpy -def numpy_fix(self, *args, **kwargs): - if self.requires_grad: - self = self.detach() - return orig_tensor_numpy(self, *args, **kwargs) - - # MPS workaround for https://github.com/pytorch/pytorch/issues/89784 -orig_cumsum = torch.cumsum -orig_Tensor_cumsum = torch.Tensor.cumsum def cumsum_fix(input, cumsum_func, *args, **kwargs): if input.device.type == 'mps': output_dtype = kwargs.get('dtype', input.dtype) @@ -199,14 +171,20 @@ def cumsum_fix(input, cumsum_func, *args, **kwargs): if has_mps(): if version.parse(torch.__version__) < version.parse("1.13"): # PyTorch 1.13 doesn't need these fixes but unfortunately is slower and has regressions that prevent training from working - torch.Tensor.to = tensor_to_fix - torch.nn.functional.layer_norm = layer_norm_fix - torch.Tensor.numpy = numpy_fix + + # MPS workaround for https://github.com/pytorch/pytorch/issues/79383 + CondFunc('torch.Tensor.to', lambda orig_func, self, *args, **kwargs: orig_func(self.contiguous(), *args, **kwargs), + lambda _, self, *args, **kwargs: self.device.type != 'mps' and (args and isinstance(args[0], torch.device) and args[0].type == 'mps' or isinstance(kwargs.get('device'), torch.device) and kwargs['device'].type == 'mps')) + # MPS workaround for https://github.com/pytorch/pytorch/issues/80800 + CondFunc('torch.nn.functional.layer_norm', lambda orig_func, *args, **kwargs: orig_func(*([args[0].contiguous()] + list(args[1:])), **kwargs), + lambda _, *args, **kwargs: args and isinstance(args[0], torch.Tensor) and args[0].device.type == 'mps') + # MPS workaround for https://github.com/pytorch/pytorch/issues/90532 + CondFunc('torch.Tensor.numpy', lambda orig_func, self, *args, **kwargs: orig_func(self.detach(), *args, **kwargs), lambda _, self, *args, **kwargs: self.requires_grad) elif version.parse(torch.__version__) > version.parse("1.13.1"): cumsum_needs_int_fix = not torch.Tensor([1,2]).to(torch.device("mps")).equal(torch.ShortTensor([1,1]).to(torch.device("mps")).cumsum(0)) cumsum_needs_bool_fix = not torch.BoolTensor([True,True]).to(device=torch.device("mps"), dtype=torch.int64).equal(torch.BoolTensor([True,False]).to(torch.device("mps")).cumsum(0)) - torch.cumsum = lambda input, *args, **kwargs: ( cumsum_fix(input, orig_cumsum, *args, **kwargs) ) - torch.Tensor.cumsum = lambda self, *args, **kwargs: ( cumsum_fix(self, orig_Tensor_cumsum, *args, **kwargs) ) - orig_narrow = torch.narrow - torch.narrow = lambda *args, **kwargs: ( orig_narrow(*args, **kwargs).clone() ) + cumsum_fix_func = lambda orig_func, input, *args, **kwargs: cumsum_fix(input, orig_func, *args, **kwargs) + CondFunc('torch.cumsum', cumsum_fix_func, None) + CondFunc('torch.Tensor.cumsum', cumsum_fix_func, None) + CondFunc('torch.narrow', lambda orig_func, *args, **kwargs: orig_func(*args, **kwargs).clone(), None)