diff --git a/SeamlessTile.py b/SeamlessTile.py index b1ac693..b1c3695 100644 --- a/SeamlessTile.py +++ b/SeamlessTile.py @@ -1,3 +1,5 @@ +import copy + import torch @@ -8,6 +10,7 @@ class SeamlessTile: "required": { "model": ("MODEL",), "tiling": (["enable", "disable"],), + "copy_model": (["Modify in place", "Make a copy"],), }, } @@ -16,14 +19,26 @@ class SeamlessTile: RETURN_TYPES = ("MODEL",) FUNCTION = "run" - def run(self, model, tiling): - for m in model.model.modules(): - if isinstance(m, torch.nn.Conv2d): - if tiling == "enable": - m.padding_mode = "circular" - else: - m.padding_mode = "zeros" - return (model,) + def run(self, model, copy_model, tiling): + if copy_model == "Modify in place": + model_copy = model + else: + model_copy = copy.deepcopy(model) + if tiling == "enable": + model_copy.model.apply(make_circular) + else: + model_copy.model.apply(unmake_circular) + return (model_copy,) + + +def make_circular(m): + if isinstance(m, torch.nn.Conv2d): + m.padding_mode = "circular" + + +def unmake_circular(m): + if isinstance(m, torch.nn.Conv2d): + m.padding_mode = "zeros" class CircularVAEDecode: @@ -51,3 +66,31 @@ class CircularVAEDecode: ]: layer.padding_mode = "zeros" return result + + +class MakeCircularVAE: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "vae": ("VAE",), + "tiling": (["enable", "disable"],), + "copy_vae": (["Modify in place", "Make a copy"],), + } + } + + RETURN_TYPES = ("VAE",) + FUNCTION = "run" + + CATEGORY = "latent" + + def run(self, vae, tiling, copy_vae): + if copy_vae == "Modify in place": + vae_copy = vae + else: + vae_copy = copy.deepcopy(vae) + if tiling == "enable": + vae_copy.first_stage_model.apply(make_circular) + else: + vae_copy.first_stage_model.apply(unmake_circular) + return (vae_copy,) diff --git a/__init__.py b/__init__.py index 84c4445..0aebc7d 100644 --- a/__init__.py +++ b/__init__.py @@ -1,13 +1,15 @@ -from .SeamlessTile import CircularVAEDecode, SeamlessTile +from .SeamlessTile import CircularVAEDecode, MakeCircularVAE, SeamlessTile NODE_CLASS_MAPPINGS = { "SeamlessTile": SeamlessTile, "CircularVAEDecode": CircularVAEDecode, + "MakeCircularVAE": MakeCircularVAE, } NODE_DISPLAY_NAME_MAPPINGS = { "SeamlessTile": "Seamless Tile", "CircularVAEDecode": "Circular VAE Decode (tile)", + "MakeCircularVAE": "Make Circular VAE", } __all__ = [NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS]