second

app.py

@@ -296,19 +296,37 @@ with gr.Blocks() as demo:
                 gradient_accumulation_steps=gr.Number(value="5", label="Gradient accumulation", interactive= True )
                 
                 add_button  = gr.Button("Run optimization")
-
-                run_optimization = partial(
-                    run_main,                 
-                    num_tokens=int(num_tokens.value),
-                    embedding_learning_rate = float(embedding_learning_rate.value), 
-                    max_emb_train_steps = int(max_emb_train_steps.value), 
-                    diffusion_model_learning_rate= float(diffusion_model_learning_rate.value), 
-                    max_diffusion_train_steps = int(max_diffusion_train_steps.value),
-                    train_batch_size=int(train_batch_size.value),
-                    gradient_accumulation_steps=int(gradient_accumulation_steps.value)
-                )
-                add_button.click(run_optimization, 
-                        inputs = [], 
+                def run_optimization_wrapper (                        
+                        num_tokens,
+                        embedding_learning_rate , 
+                        max_emb_train_steps , 
+                        diffusion_model_learning_rate , 
+                        max_diffusion_train_steps,
+                        train_batch_size,
+                        gradient_accumulation_steps
+                ):
+                    run_optimization = partial(
+                        run_main,                 
+                        num_tokens=int(num_tokens),
+                        embedding_learning_rate = float(embedding_learning_rate), 
+                        max_emb_train_steps = int(max_emb_train_steps), 
+                        diffusion_model_learning_rate= float(diffusion_model_learning_rate), 
+                        max_diffusion_train_steps = int(max_diffusion_train_steps),
+                        train_batch_size=int(train_batch_size),
+                        gradient_accumulation_steps=int(gradient_accumulation_steps)
+                    )
+                    run_optimization()
+                    
+                add_button.click(run_optimization_wrapper, 
+                        inputs = [
+                            num_tokens,
+                            embedding_learning_rate , 
+                            max_emb_train_steps , 
+                            diffusion_model_learning_rate , 
+                            max_diffusion_train_steps,
+                            train_batch_size,
+                            gradient_accumulation_steps
+                        ], 
                         outputs = []
                 )
 

pipeline_dedit_sdxl.py

@@ -0,0 +1,875 @@
+import torch
+from utils import import_model_class_from_model_name_or_path
+from transformers import AutoTokenizer
+from diffusers import (
+    AutoencoderKL,
+    DDPMScheduler,
+    StableDiffusionXLPipeline,
+    UNet2DConditionModel,
+)
+from accelerate import Accelerator
+from tqdm.auto import tqdm
+from utils import sdxl_prepare_input_decom, save_images
+import torch.nn.functional as F
+import itertools
+from peft import LoraConfig
+from controller import GroupedCAController, register_attention_disentangled_control, DummyController
+from utils import  image2latent, latent2image
+import matplotlib.pyplot as plt
+from utils_mask import  check_mask_overlap_torch
+
+device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
+max_length = 40
+class DEditSDXLPipeline:
+    def __init__(
+        self,
+        mask_list,
+        mask_label_list, 
+        mask_list_2 = None,
+        mask_label_list_2 = None, 
+        resolution = 1024,
+        num_tokens = 1
+    ):
+        super().__init__()
+        model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+        self.model_id = model_id
+        self.tokenizer       = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer", use_fast=False)
+        self.tokenizer_2     = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer_2",  use_fast=False)
+        text_encoder_cls_one = import_model_class_from_model_name_or_path(model_id, subfolder = "text_encoder")
+        text_encoder_cls_two = import_model_class_from_model_name_or_path(model_id, subfolder="text_encoder_2")
+        self.text_encoder    = text_encoder_cls_one.from_pretrained(model_id, subfolder="text_encoder" ).to(device)
+        self.text_encoder_2  = text_encoder_cls_two.from_pretrained(model_id, subfolder="text_encoder_2").to(device)
+        self.unet =  UNet2DConditionModel.from_pretrained(model_id, subfolder="unet"  )
+        self.unet.ca_dim = 2048
+        self.vae  = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix")
+        self.scheduler = DDPMScheduler.from_pretrained(model_id , subfolder="scheduler")
+        
+        self.mixed_precision = "fp16"
+        self.resolution = resolution
+        self.num_tokens = num_tokens
+        
+        self.mask_list = mask_list
+        self.mask_label_list = mask_label_list
+        notation_token_list = [phrase.split(" ")[-1] for phrase in mask_label_list]
+        placeholder_token_list = ["#"+word+"{}".format(widx) for widx, word in enumerate(notation_token_list)]
+        self.set_string_list, placeholder_token_ids = self.add_tokens(placeholder_token_list)
+        self.min_added_id = min(placeholder_token_ids)
+        self.max_added_id = max(placeholder_token_ids)
+
+        if mask_list_2 is not None:
+            self.mask_list_2 = mask_list_2
+            self.mask_label_list_2 = mask_label_list_2
+            notation_token_list_2  = [phrase.split(" ")[-1] for phrase in mask_label_list_2]
+            
+            placeholder_token_list_2 = ["$"+word+"{}".format(widx) for widx, word in enumerate(notation_token_list_2)]
+            self.set_string_list_2, placeholder_token_ids_2 = self.add_tokens(placeholder_token_list_2)
+            self.max_added_id = max(placeholder_token_ids_2)
+
+    def add_tokens_text_encoder_random_init(self, placeholder_token, num_tokens=1):
+        # Add the placeholder token in tokenizer
+        placeholder_tokens = [placeholder_token]
+        # add dummy tokens for multi-vector
+        additional_tokens = []
+        for i in range(1, num_tokens):
+            additional_tokens.append(f"{placeholder_token}_{i}")
+        placeholder_tokens += additional_tokens
+        num_added_tokens = self.tokenizer.add_tokens(placeholder_tokens) # 49408
+        num_added_tokens = self.tokenizer_2.add_tokens(placeholder_tokens) # 49408
+
+        if num_added_tokens != num_tokens:
+            raise ValueError(
+                f"The tokenizer already contains the token {placeholder_token}. Please pass a different"
+                " `placeholder_token` that is not already in the tokenizer."
+            )
+        placeholder_token_ids = self.tokenizer.convert_tokens_to_ids(placeholder_tokens)
+        placeholder_token_ids_2 = self.tokenizer_2.convert_tokens_to_ids(placeholder_tokens)  
+        assert placeholder_token_ids == placeholder_token_ids_2, "Two text encoders are expected to have same vocabs"
+        
+        self.text_encoder.resize_token_embeddings(len(self.tokenizer))
+        token_embeds = self.text_encoder.get_input_embeddings().weight.data
+        std, mean = torch.std_mean(token_embeds)
+        with torch.no_grad():
+            for token_id in placeholder_token_ids:         
+                token_embeds[token_id] = torch.randn_like(token_embeds[token_id])*std + mean
+
+        self.text_encoder_2.resize_token_embeddings(len(self.tokenizer))
+        token_embeds = self.text_encoder_2.get_input_embeddings().weight.data
+        std, mean = torch.std_mean(token_embeds)
+        with torch.no_grad():
+            for token_id in placeholder_token_ids:            
+                token_embeds[token_id] = torch.randn_like(token_embeds[token_id])*std + mean
+        
+        set_string = " ".join(self.tokenizer.convert_ids_to_tokens(placeholder_token_ids))
+
+        return set_string, placeholder_token_ids
+
+    def add_tokens(self, placeholder_token_list):
+        set_string_list = []
+        placeholder_token_ids_list = []
+        for str_idx in range(len(placeholder_token_list)):
+            placeholder_token = placeholder_token_list[str_idx]
+            set_string, placeholder_token_ids = self.add_tokens_text_encoder_random_init(placeholder_token,  num_tokens=self.num_tokens)
+            set_string_list.append(set_string)
+            placeholder_token_ids_list.append(placeholder_token_ids)
+        placeholder_token_ids = list(itertools.chain(*placeholder_token_ids_list)) 
+        return set_string_list, placeholder_token_ids
+
+    def train_emb(
+        self, 
+        image_gt,
+        set_string_list,
+        gradient_accumulation_steps = 5,
+        embedding_learning_rate = 1e-4,
+        max_emb_train_steps = 100,
+        train_batch_size = 1,
+        train_full_lora = False
+    ):
+        decom_controller =  GroupedCAController(mask_list = self.mask_list)
+        register_attention_disentangled_control(self.unet, decom_controller)
+        
+        accelerator = Accelerator(mixed_precision=self.mixed_precision, gradient_accumulation_steps=gradient_accumulation_steps)
+        self.vae.requires_grad_(False)
+        self.unet.requires_grad_(False)
+
+        self.text_encoder.requires_grad_(True)
+        self.text_encoder_2.requires_grad_(True)
+        
+        self.text_encoder.text_model.encoder.requires_grad_(False)
+        self.text_encoder.text_model.final_layer_norm.requires_grad_(False)
+        self.text_encoder.text_model.embeddings.position_embedding.requires_grad_(False)
+        
+        self.text_encoder_2.text_model.encoder.requires_grad_(False)
+        self.text_encoder_2.text_model.final_layer_norm.requires_grad_(False)
+        self.text_encoder_2.text_model.embeddings.position_embedding.requires_grad_(False)
+
+        weight_dtype = torch.float32
+        if accelerator.mixed_precision == "fp16":
+            weight_dtype = torch.float16
+        elif accelerator.mixed_precision == "bf16":
+            weight_dtype = torch.bfloat16
+
+        self.unet.to(device, dtype=weight_dtype)
+        self.vae.to(device, dtype=weight_dtype)
+
+        trainable_embmat_list_1 = [param for param in self.text_encoder.get_input_embeddings().parameters()]
+        trainable_embmat_list_2 = [param for param in self.text_encoder_2.get_input_embeddings().parameters()]
+
+        optimizer = torch.optim.AdamW(trainable_embmat_list_1 + trainable_embmat_list_2, lr=embedding_learning_rate)
+        
+        self.text_encoder, self.text_encoder_2, optimizer = accelerator.prepare(self.text_encoder, self.text_encoder_2, optimizer)
+
+        orig_embeds_params_1 = accelerator.unwrap_model(self.text_encoder)  .get_input_embeddings().weight.data.clone()
+        orig_embeds_params_2 = accelerator.unwrap_model(self.text_encoder_2).get_input_embeddings().weight.data.clone()
+
+        self.text_encoder.train()
+        self.text_encoder_2.train()
+        
+        effective_emb_train_steps = max_emb_train_steps//gradient_accumulation_steps
+        
+        if accelerator.is_main_process:
+            accelerator.init_trackers("DEdit EmbSteps", config={
+                    "embedding_learning_rate": embedding_learning_rate,
+                    "text_embedding_optimization_steps": effective_emb_train_steps,
+                })
+        global_step = 0
+        noise_scheduler = DDPMScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0" , subfolder="scheduler")
+        progress_bar = tqdm(range(0, effective_emb_train_steps), initial = global_step, desc="EmbSteps")
+        latents0 = image2latent(image_gt, vae = self.vae, dtype=weight_dtype)
+        latents0 = latents0.repeat(train_batch_size, 1, 1, 1)
+        
+        for _ in range(max_emb_train_steps):
+            with accelerator.accumulate(self.text_encoder, self.text_encoder_2):
+                latents = latents0.clone().detach()
+                noise = torch.randn_like(latents)
+                bsz = latents.shape[0]
+                timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
+                timesteps = timesteps.long()
+                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+                encoder_hidden_states_list, add_text_embeds, add_time_ids = sdxl_prepare_input_decom(
+                    set_string_list, 
+                    self.tokenizer, 
+                    self.tokenizer_2,
+                    self.text_encoder, 
+                    self.text_encoder_2, 
+                    length = max_length, 
+                    bsz = train_batch_size, 
+                    weight_dtype = weight_dtype
+                )
+
+                model_pred = self.unet(
+                    noisy_latents,
+                    timesteps, 
+                    encoder_hidden_states = encoder_hidden_states_list,
+                    cross_attention_kwargs = None,
+                    added_cond_kwargs={"text_embeds": add_text_embeds, "time_ids": add_time_ids},
+                    return_dict=False
+                )[0]
+                loss = F.mse_loss(model_pred.float(), noise.float(), reduction="mean")
+                accelerator.backward(loss)
+                optimizer.step()
+                optimizer.zero_grad()
+
+                index_no_updates = torch.ones((len(self.tokenizer),), dtype=torch.bool)
+                index_no_updates[self.min_added_id : self.max_added_id + 1] = False
+                with torch.no_grad():
+                    accelerator.unwrap_model(self.text_encoder).get_input_embeddings().weight[
+                        index_no_updates] = orig_embeds_params_1[index_no_updates]
+
+                index_no_updates = torch.ones((len(self.tokenizer_2),), dtype=torch.bool)
+                index_no_updates[self.min_added_id : self.max_added_id + 1] = False
+                with torch.no_grad():
+                    accelerator.unwrap_model(self.text_encoder_2).get_input_embeddings().weight[
+                        index_no_updates] = orig_embeds_params_2[index_no_updates]
+                        
+            logs = {"loss": loss.detach().item(), "lr": embedding_learning_rate}
+            progress_bar.set_postfix(**logs)
+            accelerator.log(logs, step=global_step)
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                global_step += 1
+                
+            if global_step >= max_emb_train_steps:
+                break
+        accelerator.wait_for_everyone()
+        accelerator.end_training()
+        self.text_encoder   =  accelerator.unwrap_model(self.text_encoder).to(dtype = weight_dtype)
+        self.text_encoder_2 =  accelerator.unwrap_model(self.text_encoder_2).to(dtype = weight_dtype)
+
+    def train_model(
+        self, 
+        image_gt, 
+        set_string_list,
+        gradient_accumulation_steps = 5,
+        max_diffusion_train_steps = 100,
+        diffusion_model_learning_rate = 1e-5,
+        train_batch_size = 1,
+        train_full_lora = False,
+        lora_rank = 4,
+        lora_alpha = 4
+    ):
+        self.unet = UNet2DConditionModel.from_pretrained(self.model_id, subfolder="unet").to(device)
+        self.unet.ca_dim = 2048
+        decom_controller =  GroupedCAController(mask_list = self.mask_list)
+        register_attention_disentangled_control(self.unet, decom_controller)
+        
+        mixed_precision = "fp16"
+        accelerator = Accelerator(gradient_accumulation_steps = gradient_accumulation_steps, mixed_precision = mixed_precision)
+        
+        weight_dtype = torch.float32
+        if accelerator.mixed_precision == "fp16":
+            weight_dtype = torch.float16
+        elif accelerator.mixed_precision == "bf16":
+            weight_dtype = torch.bfloat16
+        
+        self.vae.requires_grad_(False)
+        self.vae.to(device, dtype=weight_dtype)
+        
+        self.unet.requires_grad_(False)
+        self.unet.train()
+        
+        self.text_encoder.requires_grad_(False)
+        self.text_encoder_2.requires_grad_(False)
+        
+        if not train_full_lora:
+            trainable_params_list = []
+            for _, module in self.unet.named_modules():
+                module_name = type(module).__name__
+                if module_name == "Attention":
+                    if module.to_k.in_features == 2048: # this is cross attention:
+                        module.to_k.weight.requires_grad = True
+                        trainable_params_list.append(module.to_k.weight)
+                        if module.to_k.bias is not None:
+                            module.to_k.bias.requires_grad = True
+                            trainable_params_list.append(module.to_k.bias)                 
+                        module.to_v.weight.requires_grad = True
+                        trainable_params_list.append(module.to_v.weight)
+                        if module.to_v.bias is not None:
+                            module.to_v.bias.requires_grad = True
+                            trainable_params_list.append(module.to_v.bias)
+                        module.to_q.weight.requires_grad = True
+                        trainable_params_list.append(module.to_q.weight)
+                        if module.to_q.bias is not None:
+                            module.to_q.bias.requires_grad = True
+                            trainable_params_list.append(module.to_q.bias)
+        else:
+            unet_lora_config = LoraConfig(
+                r=lora_rank,
+                lora_alpha=lora_alpha,
+                init_lora_weights="gaussian",
+                target_modules=["to_k", "to_q", "to_v", "to_out.0"],
+            )
+            self.unet.add_adapter(unet_lora_config) 
+            print("training full parameters using lora!")
+            trainable_params_list = list(filter(lambda p: p.requires_grad, self.unet.parameters()))
+            
+        self.text_encoder.to(device, dtype=weight_dtype)
+        self.text_encoder_2.to(device, dtype=weight_dtype)         
+        optimizer = torch.optim.AdamW(trainable_params_list, lr=diffusion_model_learning_rate)                                 
+        self.unet, optimizer = accelerator.prepare(self.unet, optimizer)
+        psum2 = sum(p.numel() for p in trainable_params_list)
+
+        effective_diffusion_train_steps = max_diffusion_train_steps // gradient_accumulation_steps
+        if accelerator.is_main_process:
+            accelerator.init_trackers("textual_inversion", config={
+                    "diffusion_model_learning_rate": diffusion_model_learning_rate,
+                    "diffusion_model_optimization_steps": effective_diffusion_train_steps,
+                })
+    
+        global_step = 0
+        progress_bar = tqdm( range(0, effective_diffusion_train_steps),initial=global_step, desc="ModelSteps")
+
+        noise_scheduler = DDPMScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0" , subfolder="scheduler")
+        
+        latents0 = image2latent(image_gt, vae = self.vae, dtype=weight_dtype)
+        latents0 = latents0.repeat(train_batch_size, 1, 1, 1)
+
+        with torch.no_grad():
+            encoder_hidden_states_list, add_text_embeds, add_time_ids = sdxl_prepare_input_decom(
+                set_string_list,
+                self.tokenizer,
+                self.tokenizer_2,
+                self.text_encoder,
+                self.text_encoder_2,
+                length = max_length, 
+                bsz = train_batch_size,
+                weight_dtype = weight_dtype
+            )
+
+        for _ in range(max_diffusion_train_steps):
+            with accelerator.accumulate(self.unet):
+                latents = latents0.clone().detach()
+                noise = torch.randn_like(latents)
+                bsz = latents.shape[0]
+                timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
+                timesteps = timesteps.long()
+                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+                model_pred = self.unet(
+                    noisy_latents,
+                    timesteps, 
+                    encoder_hidden_states=encoder_hidden_states_list,
+                    cross_attention_kwargs=None, return_dict=False,
+                    added_cond_kwargs={"text_embeds": add_text_embeds, "time_ids": add_time_ids}
+                )[0]
+                loss = F.mse_loss(model_pred.float(), noise.float(), reduction="mean")
+                accelerator.backward(loss)
+                optimizer.step()
+                optimizer.zero_grad()
+
+            logs = {"loss": loss.detach().item(), "lr": diffusion_model_learning_rate}
+            progress_bar.set_postfix(**logs)
+            accelerator.log(logs, step=global_step)
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                global_step += 1
+            if global_step >=max_diffusion_train_steps:
+                break
+        accelerator.wait_for_everyone()
+        accelerator.end_training()
+        self.unet = accelerator.unwrap_model(self.unet).to(dtype = weight_dtype)
+
+    def train_emb_2imgs(
+        self,
+        image_gt_1,
+        image_gt_2, 
+        set_string_list_1,
+        set_string_list_2,
+        gradient_accumulation_steps = 5,
+        embedding_learning_rate = 1e-4,
+        max_emb_train_steps = 100,
+        train_batch_size = 1,
+        train_full_lora = False
+    ):
+        decom_controller_1 = GroupedCAController(mask_list = self.mask_list)
+        decom_controller_2 = GroupedCAController(mask_list = self.mask_list_2)
+        accelerator = Accelerator(mixed_precision=self.mixed_precision, gradient_accumulation_steps=gradient_accumulation_steps)
+        self.vae.requires_grad_(False)
+        self.unet.requires_grad_(False)
+
+        self.text_encoder.requires_grad_(True)
+        self.text_encoder_2.requires_grad_(True)
+        
+        self.text_encoder.text_model.encoder.requires_grad_(False)
+        self.text_encoder.text_model.final_layer_norm.requires_grad_(False)
+        self.text_encoder.text_model.embeddings.position_embedding.requires_grad_(False)
+        
+        self.text_encoder_2.text_model.encoder.requires_grad_(False)
+        self.text_encoder_2.text_model.final_layer_norm.requires_grad_(False)
+        self.text_encoder_2.text_model.embeddings.position_embedding.requires_grad_(False)
+
+        weight_dtype = torch.float32
+        if accelerator.mixed_precision == "fp16":
+            weight_dtype = torch.float16
+        elif accelerator.mixed_precision == "bf16":
+            weight_dtype = torch.bfloat16
+
+        self.unet.to(device, dtype=weight_dtype)
+        self.vae.to(device, dtype=weight_dtype)
+
+        
+        trainable_embmat_list_1 = [param for param in self.text_encoder.get_input_embeddings().parameters()]
+        trainable_embmat_list_2 = [param for param in self.text_encoder_2.get_input_embeddings().parameters()]
+
+        optimizer = torch.optim.AdamW(trainable_embmat_list_1 + trainable_embmat_list_2, lr=embedding_learning_rate)
+        self.text_encoder, self.text_encoder_2,  optimizer= accelerator.prepare(self.text_encoder, self.text_encoder_2,  optimizer)  ###
+        orig_embeds_params_1 = accelerator.unwrap_model(self.text_encoder)  .get_input_embeddings().weight.data.clone()
+        orig_embeds_params_2 = accelerator.unwrap_model(self.text_encoder_2).get_input_embeddings().weight.data.clone()
+
+        self.text_encoder.train()
+        self.text_encoder_2.train()
+        
+        effective_emb_train_steps = max_emb_train_steps//gradient_accumulation_steps
+        
+        if accelerator.is_main_process:
+            accelerator.init_trackers("EmbFt", config={
+                    "embedding_learning_rate": embedding_learning_rate,
+                    "text_embedding_optimization_steps": effective_emb_train_steps,
+                })
+        
+        global_step = 0
+        
+        noise_scheduler = DDPMScheduler.from_pretrained(self.model_id , subfolder="scheduler")
+        progress_bar = tqdm(range(0, effective_emb_train_steps),initial=global_step,desc="EmbSteps")
+        latents0_1 = image2latent(image_gt_1, vae = self.vae, dtype=weight_dtype)
+        latents0_1 = latents0_1.repeat(train_batch_size,1,1,1)
+        
+        latents0_2 = image2latent(image_gt_2, vae = self.vae, dtype=weight_dtype)
+        latents0_2 = latents0_2.repeat(train_batch_size,1,1,1)
+        
+        for step  in range(max_emb_train_steps):
+            with accelerator.accumulate(self.text_encoder, self.text_encoder_2):
+                latents_1 = latents0_1.clone().detach()
+                noise_1 = torch.randn_like(latents_1)
+                
+                latents_2 = latents0_2.clone().detach()
+                noise_2 = torch.randn_like(latents_2)
+                
+                bsz = latents_1.shape[0]
+
+                timesteps_1 = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents_1.device)
+                timesteps_1 = timesteps_1.long()
+                noisy_latents_1 = noise_scheduler.add_noise(latents_1, noise_1, timesteps_1)
+                
+                timesteps_2 = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents_2.device)
+                timesteps_2 = timesteps_2.long()
+                noisy_latents_2 = noise_scheduler.add_noise(latents_2, noise_2, timesteps_2)  
+                
+                register_attention_disentangled_control(self.unet, decom_controller_1)
+                encoder_hidden_states_list_1, add_text_embeds_1, add_time_ids_1 = sdxl_prepare_input_decom(
+                    set_string_list_1, 
+                    self.tokenizer,
+                    self.tokenizer_2,
+                    self.text_encoder,
+                    self.text_encoder_2, 
+                    length = max_length,
+                    bsz = train_batch_size,
+                    weight_dtype = weight_dtype
+                )
+                
+                model_pred_1 = self.unet(
+                    noisy_latents_1,
+                    timesteps_1, 
+                    encoder_hidden_states=encoder_hidden_states_list_1,
+                    cross_attention_kwargs=None,
+                    added_cond_kwargs={"text_embeds": add_text_embeds_1, "time_ids": add_time_ids_1},
+                    return_dict=False
+                )[0]
+
+                register_attention_disentangled_control(self.unet, decom_controller_2)
+                # import pdb; pdb.set_trace()
+                encoder_hidden_states_list_2, add_text_embeds_2, add_time_ids_2 = sdxl_prepare_input_decom(
+                    set_string_list_2, 
+                    self.tokenizer,
+                    self.tokenizer_2,
+                    self.text_encoder,
+                    self.text_encoder_2, 
+                    length = max_length,
+                    bsz = train_batch_size,
+                    weight_dtype = weight_dtype
+                )
+                
+                model_pred_2 = self.unet(
+                    noisy_latents_2,
+                    timesteps_2, 
+                    encoder_hidden_states = encoder_hidden_states_list_2,
+                    cross_attention_kwargs=None,
+                    added_cond_kwargs={"text_embeds": add_text_embeds_2, "time_ids": add_time_ids_2},
+                    return_dict=False
+                )[0]
+
+                loss_1 = F.mse_loss(model_pred_1.float(), noise_1.float(), reduction="mean") /2
+                loss_2 = F.mse_loss(model_pred_2.float(), noise_2.float(), reduction="mean") /2
+                loss = loss_1 + loss_2
+                accelerator.backward(loss)
+                optimizer.step()
+                optimizer.zero_grad()
+
+                index_no_updates = torch.ones((len(self.tokenizer),), dtype=torch.bool)
+                index_no_updates[self.min_added_id : self.max_added_id + 1] = False
+                with torch.no_grad():
+                    accelerator.unwrap_model(self.text_encoder).get_input_embeddings().weight[
+                        index_no_updates] = orig_embeds_params_1[index_no_updates]
+                index_no_updates = torch.ones((len(self.tokenizer_2),), dtype=torch.bool)
+                index_no_updates[self.min_added_id : self.max_added_id + 1] = False
+                with torch.no_grad():
+                    accelerator.unwrap_model(self.text_encoder_2).get_input_embeddings().weight[
+                        index_no_updates] = orig_embeds_params_2[index_no_updates]
+                        
+            logs = {"loss": loss.detach().item(), "lr": embedding_learning_rate}
+            progress_bar.set_postfix(**logs)
+            accelerator.log(logs, step=global_step)
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                global_step += 1
+                
+            if global_step >= max_emb_train_steps:
+                break
+        accelerator.wait_for_everyone()
+        accelerator.end_training()
+        self.text_encoder = accelerator.unwrap_model(self.text_encoder)  .to(dtype = weight_dtype)
+        self.text_encoder_2 = accelerator.unwrap_model(self.text_encoder_2).to(dtype = weight_dtype)
+
+    def train_model_2imgs(
+        self,
+        image_gt_1,
+        image_gt_2,
+        set_string_list_1,
+        set_string_list_2,
+        gradient_accumulation_steps = 5,
+        max_diffusion_train_steps = 100,
+        diffusion_model_learning_rate = 1e-5,
+        train_batch_size = 1,
+        train_full_lora = False,
+        lora_rank = 4,
+        lora_alpha = 4
+    ):
+        self.unet = UNet2DConditionModel.from_pretrained(self.model_id, subfolder="unet").to(device)
+        self.unet.ca_dim = 2048
+        decom_controller_1 = GroupedCAController(mask_list = self.mask_list)
+        decom_controller_2 = GroupedCAController(mask_list = self.mask_list_2)
+        
+        mixed_precision = "fp16"
+        accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps,mixed_precision=mixed_precision)
+        
+        weight_dtype = torch.float32
+        if accelerator.mixed_precision == "fp16":
+            weight_dtype = torch.float16
+        elif accelerator.mixed_precision == "bf16":
+            weight_dtype = torch.bfloat16
+        
+        
+        self.vae.requires_grad_(False)
+        self.vae.to(device, dtype=weight_dtype)
+        self.unet.requires_grad_(False)
+        self.unet.train()
+        
+        self.text_encoder.requires_grad_(False)
+        self.text_encoder_2.requires_grad_(False)
+        if not train_full_lora:
+            trainable_params_list = []
+            for name, module in self.unet.named_modules():
+                module_name = type(module).__name__
+                if module_name == "Attention":
+                    if module.to_k.in_features == 2048: # this is cross attention:
+                        module.to_k.weight.requires_grad = True
+                        trainable_params_list.append(module.to_k.weight)
+                        if module.to_k.bias is not None:
+                            module.to_k.bias.requires_grad = True
+                            trainable_params_list.append(module.to_k.bias)
+                            
+                        module.to_v.weight.requires_grad = True
+                        trainable_params_list.append(module.to_v.weight)
+                        if module.to_v.bias is not None:
+                            module.to_v.bias.requires_grad = True
+                            trainable_params_list.append(module.to_v.bias)
+                        module.to_q.weight.requires_grad = True
+                        trainable_params_list.append(module.to_q.weight)
+                        if module.to_q.bias is not None:
+                            module.to_q.bias.requires_grad = True
+                            trainable_params_list.append(module.to_q.bias)
+        else:
+            unet_lora_config = LoraConfig(
+                r = lora_rank,
+                lora_alpha = lora_alpha,
+                init_lora_weights="gaussian",
+                target_modules=["to_k", "to_q", "to_v", "to_out.0"],
+            )
+            self.unet.add_adapter(unet_lora_config) 
+            print("training full parameters using lora!")
+            trainable_params_list = list(filter(lambda p: p.requires_grad, self.unet.parameters()))
+            
+        self.text_encoder.to(device, dtype=weight_dtype)
+        self.text_encoder_2.to(device, dtype=weight_dtype)         
+        optimizer = torch.optim.AdamW(trainable_params_list, lr=diffusion_model_learning_rate)                                 
+        self.unet, optimizer = accelerator.prepare(self.unet, optimizer)
+        psum2 = sum(p.numel() for p in trainable_params_list)
+
+        effective_diffusion_train_steps = max_diffusion_train_steps // gradient_accumulation_steps
+        if accelerator.is_main_process:
+            accelerator.init_trackers("ModelFt", config={
+                    "diffusion_model_learning_rate": diffusion_model_learning_rate,
+                    "diffusion_model_optimization_steps": effective_diffusion_train_steps,
+                })
+    
+        global_step = 0
+        progress_bar = tqdm(range(0, effective_diffusion_train_steps),initial=global_step, desc="ModelSteps")
+        noise_scheduler = DDPMScheduler.from_pretrained(self.model_id, subfolder="scheduler")
+        
+        latents0_1 = image2latent(image_gt_1, vae = self.vae, dtype=weight_dtype)
+        latents0_1 = latents0_1.repeat(train_batch_size, 1, 1, 1)
+
+        latents0_2 = image2latent(image_gt_2, vae = self.vae, dtype=weight_dtype)
+        latents0_2 = latents0_2.repeat(train_batch_size,1, 1, 1)
+        
+        with torch.no_grad():
+            encoder_hidden_states_list_1, add_text_embeds_1, add_time_ids_1 = sdxl_prepare_input_decom(
+                set_string_list_1,
+                self.tokenizer,
+                self.tokenizer_2,
+                self.text_encoder,
+                self.text_encoder_2,
+                length = max_length, 
+                bsz = train_batch_size, 
+                weight_dtype = weight_dtype
+            )
+            encoder_hidden_states_list_2, add_text_embeds_2, add_time_ids_2 = sdxl_prepare_input_decom(
+                set_string_list_2,
+                self.tokenizer,
+                self.tokenizer_2,
+                self.text_encoder,
+                self.text_encoder_2,
+                length = max_length, 
+                bsz = train_batch_size,
+                weight_dtype = weight_dtype
+            )
+            
+        for _ in range(max_diffusion_train_steps):
+            with accelerator.accumulate(self.unet):
+                latents_1 = latents0_1.clone().detach()
+                noise_1 = torch.randn_like(latents_1)
+                bsz = latents_1.shape[0]
+                timesteps_1 = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents_1.device)
+                timesteps_1 = timesteps_1.long()
+                noisy_latents_1 = noise_scheduler.add_noise(latents_1, noise_1, timesteps_1)
+                
+                latents_2 = latents0_2.clone().detach()
+                noise_2 = torch.randn_like(latents_2)
+                bsz = latents_2.shape[0]
+                timesteps_2 = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents_2.device)
+                timesteps_2 = timesteps_2.long()
+                noisy_latents_2 = noise_scheduler.add_noise(latents_2, noise_2, timesteps_2)
+                
+                register_attention_disentangled_control(self.unet, decom_controller_1)
+                model_pred_1 = self.unet(
+                    noisy_latents_1,
+                    timesteps_1, 
+                    encoder_hidden_states = encoder_hidden_states_list_1,
+                    cross_attention_kwargs = None,
+                    return_dict = False,
+                    added_cond_kwargs = {"text_embeds": add_text_embeds_1, "time_ids": add_time_ids_1}
+                )[0]
+                
+                register_attention_disentangled_control(self.unet, decom_controller_2)
+                model_pred_2 = self.unet(
+                    noisy_latents_2,
+                    timesteps_2, 
+                    encoder_hidden_states = encoder_hidden_states_list_2,
+                    cross_attention_kwargs = None,
+                    return_dict=False,
+                    added_cond_kwargs={"text_embeds": add_text_embeds_2, "time_ids": add_time_ids_2}
+                )[0]
+                
+                loss_1 = F.mse_loss(model_pred_1.float(), noise_1.float(), reduction="mean")
+                loss_2 = F.mse_loss(model_pred_2.float(), noise_2.float(), reduction="mean")
+                loss = loss_1 + loss_2
+                accelerator.backward(loss)
+                optimizer.step()
+                optimizer.zero_grad()
+                
+
+            logs = {"loss": loss.detach().item(), "lr": diffusion_model_learning_rate}
+            progress_bar.set_postfix(**logs)
+            accelerator.log(logs, step=global_step)
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                global_step += 1
+                
+            if global_step >=max_diffusion_train_steps:
+                break
+        accelerator.wait_for_everyone()
+        accelerator.end_training()
+        self.unet = accelerator.unwrap_model(self.unet).to(dtype = weight_dtype)
+
+    @torch.no_grad()
+    def backward_zT_to_z0_euler_decom(
+        self,
+        zT,
+        cond_emb_list,
+        cond_add_text_embeds,
+        add_time_ids,
+        uncond_emb=None,
+        guidance_scale = 1,
+        num_sampling_steps = 20,
+        cond_controller = None,
+        uncond_controller = None,
+        mask_hard = None, 
+        mask_soft = None,
+        orig_image = None,
+        return_intermediate = False,
+        strength = 1
+    ):
+        latent_cur = zT
+        if uncond_emb is None:
+            uncond_emb = torch.zeros(zT.shape[0], 77, 2048).to(dtype = zT.dtype, device = zT.device)
+            uncond_add_text_embeds = torch.zeros(1, 1280).to(dtype = zT.dtype, device = zT.device)
+        if mask_soft is not None:
+            init_latents_orig = image2latent(orig_image, self.vae, dtype=self.vae.dtype)
+            length = init_latents_orig.shape[-1]
+            noise = torch.randn_like(init_latents_orig)
+            mask_soft = torch.nn.functional.interpolate(mask_soft.float().unsqueeze(0).unsqueeze(0), (length, length)).to(self.vae.dtype) ###
+        if mask_hard is not None:
+            init_latents_orig = image2latent(orig_image, self.vae, dtype=self.vae.dtype)
+            length = init_latents_orig.shape[-1]
+            noise = torch.randn_like(init_latents_orig) 
+            mask_hard = torch.nn.functional.interpolate(mask_hard.float().unsqueeze(0).unsqueeze(0), (length, length)).to(self.vae.dtype) ###
+
+        intermediate_list = [latent_cur.detach()]
+        for i in tqdm(range(num_sampling_steps)):
+            t = self.scheduler.timesteps[i]
+            latent_input = self.scheduler.scale_model_input(latent_cur, t)
+            
+            register_attention_disentangled_control(self.unet, uncond_controller)
+            noise_pred_uncond = self.unet(latent_input, t, 
+                            encoder_hidden_states=uncond_emb,
+                            added_cond_kwargs={"text_embeds": uncond_add_text_embeds, "time_ids": add_time_ids},
+                            return_dict=False,)[0]
+            
+            register_attention_disentangled_control(self.unet, cond_controller)
+            noise_pred_cond = self.unet(latent_input, t, 
+                    encoder_hidden_states=cond_emb_list,
+                    added_cond_kwargs={"text_embeds": cond_add_text_embeds, "time_ids": add_time_ids},
+                    return_dict=False,)[0]
+            
+            noise_pred =  noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+            latent_cur = self.scheduler.step(noise_pred, t, latent_cur, generator = None, return_dict=False)[0]
+            if return_intermediate is True:
+                intermediate_list.append(latent_cur)
+            if mask_hard is not None and mask_soft is not None and i <= strength *num_sampling_steps:
+                init_latents_proper = self.scheduler.add_noise(init_latents_orig, noise, torch.tensor([t]))
+                mask = mask_soft.to(latent_cur.device, latent_cur.dtype) + mask_hard.to(latent_cur.device, latent_cur.dtype)
+                latent_cur = (init_latents_proper * mask) + (latent_cur * (1 - mask))
+
+            elif mask_hard is not None and mask_soft is not None and i > strength *num_sampling_steps:
+                init_latents_proper = self.scheduler.add_noise(init_latents_orig, noise, torch.tensor([t]))
+                mask = mask_hard.to(latent_cur.device, latent_cur.dtype)
+                latent_cur = (init_latents_proper * mask) + (latent_cur * (1 - mask))
+
+            elif mask_hard is None and mask_soft is not None and i <= strength *num_sampling_steps:
+                init_latents_proper = self.scheduler.add_noise(init_latents_orig, noise, torch.tensor([t]))
+                mask = mask_soft.to(latent_cur.device, latent_cur.dtype)
+                latent_cur = (init_latents_proper * mask) + (latent_cur * (1 - mask))
+                
+            elif mask_hard is None and mask_soft is not None and i > strength *num_sampling_steps:
+                pass
+            
+            elif mask_hard is not None and mask_soft is None:
+                init_latents_proper = self.scheduler.add_noise(init_latents_orig, noise, torch.tensor([t]))
+                mask = mask_hard.to(latent_cur.dtype)
+                latent_cur = (init_latents_proper * mask) + (latent_cur * (1 - mask))
+                
+            else: # hard and soft are both none
+                pass
+            
+        if return_intermediate is True:
+            return latent_cur, intermediate_list
+        else:
+            return latent_cur
+        
+    @torch.no_grad()
+    def sampling(
+        self,
+        set_string_list, 
+        cond_controller = None, 
+        uncond_controller = None,
+        guidance_scale = 7,
+        num_sampling_steps = 20,
+        mask_hard = None,
+        mask_soft = None,
+        orig_image = None,
+        strength = 1.,
+        num_imgs = 1,
+        normal_token_id_list = [],
+        seed = 1
+    ):
+        weight_dtype = torch.float16
+        self.scheduler.set_timesteps(num_sampling_steps)
+        self.unet.to(device, dtype=weight_dtype)
+        self.vae.to(device, dtype=weight_dtype)
+        self.text_encoder.to(device, dtype=weight_dtype)
+        self.text_encoder_2.to(device, dtype=weight_dtype)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+
+        vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        zT = torch.randn(num_imgs, 4, self.resolution//vae_scale_factor,self.resolution//vae_scale_factor).to(device,dtype=weight_dtype)
+        zT = zT * self.scheduler.init_noise_sigma   
+        
+        cond_emb_list, cond_add_text_embeds, add_time_ids = sdxl_prepare_input_decom(
+            set_string_list, 
+            self.tokenizer,
+            self.tokenizer_2,
+            self.text_encoder,
+            self.text_encoder_2, 
+            length = max_length,
+            bsz = num_imgs,
+            weight_dtype = weight_dtype,
+            normal_token_id_list = normal_token_id_list
+        )
+
+        z0 = self.backward_zT_to_z0_euler_decom(zT, cond_emb_list, cond_add_text_embeds, add_time_ids,
+                guidance_scale = guidance_scale, num_sampling_steps = num_sampling_steps,
+                cond_controller = cond_controller, uncond_controller = uncond_controller, 
+                mask_hard = mask_hard, mask_soft = mask_soft,  orig_image =orig_image, strength = strength
+             )
+        x0 = latent2image(z0, vae = self.vae)
+        return x0
+    
+    @torch.no_grad()                 
+    def inference_with_mask(
+        self,
+        save_path,
+        guidance_scale = 3,
+        num_sampling_steps = 50,
+        strength = 1,
+        mask_soft = None,
+        mask_hard= None,
+        orig_image=None,
+        mask_list = None,
+        num_imgs = 1,
+        seed = 1,
+        set_string_list = None
+    ):
+        if mask_list is not None:
+            mask_list = [m.to(device) for m in mask_list]
+        else:
+            mask_list = self.mask_list
+        if set_string_list is not None:
+            self.set_string_list = set_string_list
+
+        if mask_hard is not None and mask_soft is not None:
+            check_mask_overlap_torch(mask_hard, mask_soft)
+        null_controller = DummyController()
+        decom_controller = GroupedCAController(mask_list = mask_list)
+        x0 = self.sampling(
+            self.set_string_list, 
+            guidance_scale = guidance_scale,
+            num_sampling_steps = num_sampling_steps,
+            strength = strength,
+            cond_controller = decom_controller,
+            uncond_controller = null_controller, 
+            mask_soft = mask_soft,
+            mask_hard = mask_hard,
+            orig_image = orig_image,
+            num_imgs = num_imgs,
+            seed = seed
+        )
+        save_images(x0, save_path)