Pixel shuffle adapter. Some bug fixes thrown in

toolkit/models/base_model.py

@@ -1,5 +1,6 @@
 import copy
 import gc
+import inspect
 import json
 import random
 import shutil
@@ -230,6 +231,20 @@ class BaseModel:
     def is_lumina2(self):
         return self.arch == 'lumina2'
 
+    def get_bucket_divisibility(self):
+        if self.vae is None:
+            return 8
+        try:
+            divisibility = 2 ** (len(self.vae.config['block_out_channels']) - 1)
+        except:
+            # if we have a custom vae, it might not have this
+            divisibility = 8
+        
+        # flux packs this again,
+        if self.is_flux:
+            divisibility = divisibility * 4
+        return divisibility
+
     # these must be implemented in child classes
     def load_model(self):
         # override this in child classes
@@ -797,13 +812,20 @@ class BaseModel:
             self.unet.to(self.device_torch)
         if self.unet.dtype != self.torch_dtype:
             self.unet = self.unet.to(dtype=self.torch_dtype)
+            
+        # check if get_noise prediction has guidance_embedding_scale
+        # if it does not, we dont pass it
+        signatures =  inspect.signature(self.get_noise_prediction).parameters
+        
+        if 'guidance_embedding_scale' in signatures:
+            kwargs['guidance_embedding_scale'] = guidance_embedding_scale
+        if 'bypass_guidance_embedding' in signatures:
+            kwargs['bypass_guidance_embedding'] = bypass_guidance_embedding
 
         noise_pred = self.get_noise_prediction(
             latent_model_input=latent_model_input,
             timestep=timestep,
             text_embeddings=text_embeddings,
-            guidance_embedding_scale=guidance_embedding_scale,
-            bypass_guidance_embedding=bypass_guidance_embedding,
             **kwargs
         )
 

toolkit/models/subpixel_adapter.py

@@ -0,0 +1,302 @@
+import inspect
+import weakref
+import torch
+from typing import TYPE_CHECKING
+from toolkit.lora_special import LoRASpecialNetwork
+from diffusers import FluxTransformer2DModel
+# weakref
+from toolkit.pixel_shuffle_encoder import AutoencoderPixelMixer
+
+
+if TYPE_CHECKING:
+    from toolkit.stable_diffusion_model import StableDiffusion
+    from toolkit.config_modules import AdapterConfig, TrainConfig, ModelConfig
+    from toolkit.custom_adapter import CustomAdapter
+    
+
+
+class InOutModule(torch.nn.Module):
+    def __init__(
+        self,
+        adapter: 'SubpixelAdapter',
+        orig_layer: torch.nn.Linear,
+        in_channels=64,
+        out_channels=3072
+    ):
+        super().__init__()
+        # only do the weight for the new input. We combine with the original linear layer
+        self.x_embedder = torch.nn.Linear(
+            in_channels,
+            out_channels,
+            bias=True,
+        )
+        
+        self.proj_out = torch.nn.Linear(
+            out_channels,
+            in_channels,
+            bias=True,
+        )
+        # make sure the weight is float32
+        self.x_embedder.weight.data = self.x_embedder.weight.data.float()
+        self.x_embedder.bias.data = self.x_embedder.bias.data.float()
+        
+        self.proj_out.weight.data = self.proj_out.weight.data.float()
+        self.proj_out.bias.data = self.proj_out.bias.data.float()
+        
+        self.adapter_ref: weakref.ref = weakref.ref(adapter)
+        self.orig_layer_ref: weakref.ref = weakref.ref(orig_layer)
+        
+    @classmethod
+    def from_model(
+        cls, 
+        model: FluxTransformer2DModel, 
+        adapter: 'SubpixelAdapter',
+        num_channels: int = 768,
+        downscale_factor: int = 8
+    ):
+        if model.__class__.__name__ == 'FluxTransformer2DModel':            
+            
+            x_embedder: torch.nn.Linear = model.x_embedder
+            proj_out: torch.nn.Linear = model.proj_out
+            in_out_module = cls(
+                adapter, 
+                orig_layer=x_embedder,
+                in_channels=num_channels,
+                out_channels=x_embedder.out_features,
+            )
+            
+            # hijack the forward method
+            x_embedder._orig_ctrl_lora_forward = x_embedder.forward
+            x_embedder.forward = in_out_module.in_forward
+            proj_out._orig_ctrl_lora_forward = proj_out.forward
+            proj_out.forward = in_out_module.out_forward
+
+            # update the config of the transformer
+            model.config.in_channels = num_channels
+            model.config["in_channels"] = num_channels
+            model.config.out_channels = num_channels
+            model.config["out_channels"] = num_channels
+            
+            # replace the vae of the model
+            sd = adapter.sd_ref()
+            sd.vae = AutoencoderPixelMixer(
+                in_channels=3,
+                downscale_factor=downscale_factor
+            )
+            
+            sd.pipeline.vae = sd.vae
+            
+            return in_out_module
+        else:
+            raise ValueError("Model not supported") 
+        
+    @property
+    def is_active(self):
+        return self.adapter_ref().is_active
+        
+    
+    def in_forward(self, x):
+        if not self.is_active:
+            # make sure lora is not active
+            if self.adapter_ref().control_lora is not None:
+                self.adapter_ref().control_lora.is_active = False
+            return self.orig_layer_ref()._orig_ctrl_lora_forward(x)
+        
+        # make sure lora is active
+        if self.adapter_ref().control_lora is not None:
+            self.adapter_ref().control_lora.is_active = True
+        
+        orig_device = x.device
+        orig_dtype = x.dtype
+    
+        x = x.to(self.x_embedder.weight.device, dtype=self.x_embedder.weight.dtype)
+        
+        x = self.x_embedder(x)
+        
+        x = x.to(orig_device, dtype=orig_dtype)
+        return x
+    
+    def out_forward(self, x):
+        if not self.is_active:
+            # make sure lora is not active
+            if self.adapter_ref().control_lora is not None:
+                self.adapter_ref().control_lora.is_active = False
+            return self.orig_layer_ref()._orig_ctrl_lora_forward(x)
+        
+        # make sure lora is active
+        if self.adapter_ref().control_lora is not None:
+            self.adapter_ref().control_lora.is_active = True
+        
+        orig_device = x.device
+        orig_dtype = x.dtype
+    
+        x = x.to(self.proj_out.weight.device, dtype=self.proj_out.weight.dtype)
+        
+        x = self.proj_out(x)
+        
+        x = x.to(orig_device, dtype=orig_dtype)
+        return x
+    
+
+
+class SubpixelAdapter(torch.nn.Module):
+    def __init__(
+        self,
+        adapter: 'CustomAdapter',
+        sd: 'StableDiffusion',
+        config: 'AdapterConfig',
+        train_config: 'TrainConfig'
+    ):
+        super().__init__()
+        self.adapter_ref: weakref.ref = weakref.ref(adapter)
+        self.sd_ref = weakref.ref(sd)
+        self.model_config: ModelConfig = sd.model_config
+        self.network_config = config.lora_config
+        self.train_config = train_config
+        self.device_torch = sd.device_torch
+        self.control_lora = None
+        
+        if self.network_config is not None:
+        
+            network_kwargs = {} if self.network_config.network_kwargs is None else self.network_config.network_kwargs
+            if hasattr(sd, 'target_lora_modules'):
+                network_kwargs['target_lin_modules'] = self.sd.target_lora_modules
+                
+            if 'ignore_if_contains' not in network_kwargs:
+                network_kwargs['ignore_if_contains'] = []
+            
+            # always ignore x_embedder
+            network_kwargs['ignore_if_contains'].append('transformer.x_embedder')
+            network_kwargs['ignore_if_contains'].append('transformer.proj_out')
+                
+            self.control_lora = LoRASpecialNetwork(
+                text_encoder=sd.text_encoder,
+                unet=sd.unet,
+                lora_dim=self.network_config.linear,
+                multiplier=1.0,
+                alpha=self.network_config.linear_alpha,
+                train_unet=self.train_config.train_unet,
+                train_text_encoder=self.train_config.train_text_encoder,
+                conv_lora_dim=self.network_config.conv,
+                conv_alpha=self.network_config.conv_alpha,
+                is_sdxl=self.model_config.is_xl or self.model_config.is_ssd,
+                is_v2=self.model_config.is_v2,
+                is_v3=self.model_config.is_v3,
+                is_pixart=self.model_config.is_pixart,
+                is_auraflow=self.model_config.is_auraflow,
+                is_flux=self.model_config.is_flux,
+                is_lumina2=self.model_config.is_lumina2,
+                is_ssd=self.model_config.is_ssd,
+                is_vega=self.model_config.is_vega,
+                dropout=self.network_config.dropout,
+                use_text_encoder_1=self.model_config.use_text_encoder_1,
+                use_text_encoder_2=self.model_config.use_text_encoder_2,
+                use_bias=False,
+                is_lorm=False,
+                network_config=self.network_config,
+                network_type=self.network_config.type,
+                transformer_only=self.network_config.transformer_only,
+                is_transformer=sd.is_transformer,
+                base_model=sd,
+                **network_kwargs
+            )
+            self.control_lora.force_to(self.device_torch, dtype=torch.float32)
+            self.control_lora._update_torch_multiplier()
+            self.control_lora.apply_to(
+                sd.text_encoder,
+                sd.unet,
+                self.train_config.train_text_encoder,
+                self.train_config.train_unet
+            )
+            self.control_lora.can_merge_in = False
+            self.control_lora.prepare_grad_etc(sd.text_encoder, sd.unet)
+            if self.train_config.gradient_checkpointing:
+                self.control_lora.enable_gradient_checkpointing()
+        
+        downscale_factor = config.subpixel_downscale_factor
+        if downscale_factor == 8:
+            num_channels = 768
+        elif downscale_factor == 16:
+            num_channels = 3072
+        else:
+            raise ValueError(
+                f"downscale_factor {downscale_factor} not supported"
+            )
+        
+        self.in_out: InOutModule = InOutModule.from_model(
+            sd.unet_unwrapped, 
+            self,
+            num_channels=num_channels, # packed channels
+            downscale_factor=downscale_factor
+        )
+        self.in_out.to(self.device_torch)
+
+    def get_params(self):
+        if self.control_lora is not None:
+            config = {
+                'text_encoder_lr': self.train_config.lr,
+                'unet_lr': self.train_config.lr,
+            }
+            sig = inspect.signature(self.control_lora.prepare_optimizer_params)
+            if 'default_lr' in sig.parameters:
+                config['default_lr'] = self.train_config.lr
+            if 'learning_rate' in sig.parameters:
+                config['learning_rate'] = self.train_config.lr
+            params_net = self.control_lora.prepare_optimizer_params(
+                **config
+            )
+            
+            # we want only tensors here
+            params = []
+            for p in params_net:
+                if isinstance(p, dict):
+                    params += p["params"]
+                elif isinstance(p, torch.Tensor):
+                    params.append(p)
+                elif isinstance(p, list):
+                    params += p
+        else:
+            params = []
+            
+        # make sure the embedder is float32
+        self.in_out.to(torch.float32)
+        
+        params += list(self.in_out.parameters())
+        
+        # we need to be able to yield from the list like yield from params
+
+        return params
+    
+    def load_weights(self, state_dict, strict=True):
+        lora_sd = {}
+        img_embedder_sd = {}
+        for key, value in state_dict.items():
+            if "transformer.x_embedder" in key:
+                new_key = key.replace("transformer.", "")
+                img_embedder_sd[new_key] = value
+            elif "transformer.proj_out" in key:
+                new_key = key.replace("transformer.", "")
+                img_embedder_sd[new_key] = value
+            else:
+                lora_sd[key] = value
+        
+        # todo process state dict before loading
+        if self.control_lora is not None:
+            self.control_lora.load_weights(lora_sd)
+        # automatically upgrade the x imbedder if more dims are added
+        self.in_out.load_state_dict(img_embedder_sd, strict=False)
+        
+    def get_state_dict(self):
+        if self.control_lora is not None:
+            lora_sd = self.control_lora.get_state_dict(dtype=torch.float32)
+        else:
+            lora_sd = {}
+        # todo make sure we match loras elseware. 
+        img_embedder_sd = self.in_out.state_dict()
+        for key, value in img_embedder_sd.items():
+            lora_sd[f"transformer.{key}"] = value
+        return lora_sd
+    
+    @property
+    def is_active(self):
+        return self.adapter_ref().is_active

toolkit/models/wan21/wan21.py

@@ -318,6 +318,9 @@ class Wan21(BaseModel):
 
         # cache for holding noise
         self.effective_noise = None
+        
+    def get_bucket_divisibility(self):
+        return 16
 
     # static method to get the scheduler
     @staticmethod