diff --git a/nodes.py b/nodes.py index 5a51473..4958baf 100644 --- a/nodes.py +++ b/nodes.py @@ -170,6 +170,26 @@ class TweakStyleGANLatents: return (z, ) +class StyleGANLatentFromBatch: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "stylegan_latent": ("STYLEGAN_LATENT", ), + "index": ("INT", {"default": 0, "min": 0}), + }, + } + + RETURN_TYPES = ("STYLEGAN_LATENT",) + FUNCTION = "generate_latent" + CATEGORY = "StyleGAN/extra" + + def generate_latent(self, stylegan_latent, index): + clipped_index = min(index, stylegan_latent.size(0) - 1) + z = stylegan_latent[clipped_index].unsqueeze(0).detach().clone() + + return (z, ) + NODE_CLASS_MAPPINGS = { "LoadStyleGAN": LoadStyleGAN, "GenerateStyleGANLatent": GenerateStyleGANLatent, @@ -177,6 +197,7 @@ NODE_CLASS_MAPPINGS = { "BlendStyleGANLatents": BlendStyleGANLatents, "BatchAverageStyleGANLatents": BatchAverageStyleGANLatents, "TweakStyleGANLatents": TweakStyleGANLatents, + "StyleGANLatentFromBatch": StyleGANLatentFromBatch, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -186,4 +207,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "BlendStyleGANLatents": "Blend StyleGAN Latents (lerp or slerp)", "BatchAverageStyleGANLatents": "Batch Average StyleGAN Latents", "TweakStyleGANLatents": "Tweak StyleGAN Latents", + "StyleGANLatentFromBatch": "StyleGAN Latent From Batch", } \ No newline at end of file