repo init

diffusion_pipeline/gemma.py

@@ -0,0 +1,53 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoTokenizer, Gemma2ForTokenClassification, BitsAndBytesConfig
+
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+torch.set_float32_matmul_precision("high")
+
+def repeat_function(xs, max_length = 128):
+    new_xs = []
+    for x in xs:
+        if x.shape[1] >= max_length-1:
+            new_xs.append(x[:,:max_length-1,:])
+        else:
+            new_xs.append(x)
+    xs = new_xs
+    mean_xs = [x.mean(1,keepdim=True).expand(-1,max_length - x.shape[1],-1) for x in xs]
+    xs = [torch.cat([x,mean_x],1) for mean_x, x in zip(mean_xs, xs)]
+    return xs
+
+class Gemma2Model(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b", )
+        self.tokenizer_max_length = 128
+        # quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+
+        self.model = Gemma2ForTokenClassification.from_pretrained(
+            "google/gemma-2-2b",
+            # device_map="auto",
+            # quantization_config=quantization_config,
+        ).float()
+        self.model.score = nn.Identity()
+    
+    @torch.no_grad()
+    def forward(self, input_prompt):
+        input_prompt = list(input_prompt)
+        outputs = []
+        for _input_prompt in input_prompt:
+            input_ids = self.tokenizer(_input_prompt, add_special_tokens=False, max_length=77, return_tensors="pt").to("cuda")
+            _outputs = self.model(**input_ids)["logits"]
+            outputs.append(_outputs)
+        outputs = repeat_function(outputs)
+        outputs = torch.cat(outputs,0)
+        return outputs
+
+if __name__ == "__main__":
+    model = Gemma2Model().cuda()
+    input_text = ["Write me a poem about Machine Learning.", "Write me a poem about Deep Learning."]
+    print(model(input_text))
+    print(model(input_text)[0].shape)
+    print(model(input_text).shape)

diffusion_pipeline/lora.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn as nn
+
+class LoRALayer(torch.nn.Module):
+    def __init__(self, in_dim, out_dim, rank, alpha):
+        super().__init__()
+        std_dev = 1 / torch.sqrt(torch.tensor(rank).float())
+        self.A = torch.nn.Parameter(torch.randn(in_dim, rank) * std_dev)
+        self.B = torch.nn.Parameter(torch.zeros(rank, out_dim))
+        self.alpha = alpha
+
+    def forward(self, x):
+        x = self.alpha * (x @ self.A @ self.B)
+        return x
+
+class LinearWithLoRA(torch.nn.Module):
+    def __init__(self, linear, rank, alpha, 
+                 weak_lora_alpha=0.1, number_of_lora=1):
+        super().__init__()
+        self.linear = linear
+        self.lora = nn.ModuleList([LoRALayer(
+            linear.in_features, linear.out_features, rank, alpha
+        ) for _ in range(number_of_lora)])
+        self.use_lora = True
+        self.lora_idx = 0
+        
+    def forward(self, x):
+        if self.use_lora:
+            return self.linear(x) + self.lora[self.lora_idx](x)
+        else:
+            return self.linear(x)
+
+def replace_linear_with_lora(module, rank=64, alpha=1., tag=0, weak_lora_alpha=0.1, number_of_lora=1):
+    for name, child in module.named_children():
+        if isinstance(child, nn.Linear):
+            setattr(module, name, LinearWithLoRA(child, rank, alpha, weak_lora_alpha=weak_lora_alpha, number_of_lora=number_of_lora))
+        else:
+            replace_linear_with_lora(child, rank, alpha, tag, weak_lora_alpha=weak_lora_alpha, number_of_lora=number_of_lora)
+
+
+def lora_false(model, lora_idx=0):
+    for name, module in model.named_modules():
+        if isinstance(module, LinearWithLoRA):
+            module.use_lora = False
+            module.lora_idx = lora_idx
+
+def lora_true(model, lora_idx=0):
+    for name, module in model.named_modules():
+        if isinstance(module, LinearWithLoRA):
+            module.use_lora = True
+            module.lora_idx = lora_idx
+            for i, lora in enumerate(module.lora):
+                if i != lora_idx:
+                    lora.A.requires_grad = False
+                    lora.B.requires_grad = False
+                    if lora.A.grad is not None:
+                        del lora.A.grad
+                    if lora.B.grad is not None:
+                        del lora.B.grad
+                else:
+                    lora.A.requires_grad = True
+                    lora.B.requires_grad = True

diffusion_pipeline/refine_model.py

@@ -0,0 +1,526 @@
+
+import torch
+import torch.nn as nn
+import os
+import json
+import torch.nn.functional as F
+import random
+from torch.utils.data import Dataset
+from transformers import AutoTokenizer
+from glob import glob
+import math
+from PIL import Image
+device = torch.device('cuda')
+import numpy as np
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from diffusers.utils import logging
+from diffusers.models.embeddings import PatchEmbed
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.attention import BasicTransformerBlock
+from diffusers.models.normalization import AdaLayerNormContinuous
+from torchvision import transforms
+    
+def add_hook_to_module(model, module_name):
+    outputs = []
+    def hook(module, input, output):
+        outputs.append(output)
+    module = dict(model.named_modules()).get(module_name)
+    if module is None:
+        raise ValueError(f"can't find module {module_name}")
+    hook_handle = module.register_forward_hook(hook)
+    return hook_handle, outputs
+
+class PromptSD35Net(nn.Module):
+
+    def __init__(self,         
+        sample_size: int = 128,
+        patch_size: int = 2,
+        in_channels: int = 16,
+        num_layers: int = 8,
+        attention_head_dim: int = 64,
+        num_attention_heads: int = 24,
+        out_channels: int = 16,
+        pos_embed_max_size: int = 192
+        ):
+        super().__init__()
+        self.sample_size = sample_size
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.num_layers = num_layers
+        self.attention_head_dim = attention_head_dim
+        self.num_attention_heads = num_attention_heads
+        self.out_channels = out_channels
+        self.pos_embed_max_size = pos_embed_max_size
+        self.inner_dim = self.num_attention_heads * self.attention_head_dim
+        
+        self.pos_embed = PatchEmbed(
+            height=self.sample_size,
+            width=self.sample_size,
+            patch_size=self.patch_size,
+            in_channels=self.in_channels,
+            embed_dim=self.inner_dim,
+            pos_embed_max_size=pos_embed_max_size
+        )
+        
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.num_attention_heads,
+                    attention_head_dim=self.attention_head_dim,
+                    ff_inner_dim=2*self.inner_dim   # mult should be 4 by default
+                )
+                for i in range(self.num_layers)
+            ]
+        )
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+        
+        self.noise_shape = (1, 16, 128, 128) # (667, 4096)
+        self.pre8_linear = nn.Sequential(nn.Linear(4096, 128), nn.SiLU(), nn.LayerNorm(128), nn.Linear(128, 1536))
+        self.pre16_linear = nn.Sequential(nn.Linear(4096, 128), nn.SiLU(), nn.LayerNorm(128), nn.Linear(128, 1536))
+        self.pre24_linear = nn.Sequential(nn.Linear(4096, 128), nn.SiLU(), nn.LayerNorm(128), nn.Linear(128, 1536))
+
+        self.pre8_linear2 = nn.Sequential(nn.Linear(4096, 128), nn.SiLU(), nn.LayerNorm(128), nn.Linear(128, 1536))
+        self.pre16_linear2 = nn.Sequential(nn.Linear(4096, 128), nn.SiLU(), nn.LayerNorm(128), nn.Linear(128, 1536))
+        self.pre24_linear2 = nn.Sequential(nn.Linear(4096, 128), nn.SiLU(), nn.LayerNorm(128), nn.Linear(128, 1536))
+        
+        self.last_linear = nn.Sequential(nn.Linear(4096, 128), nn.SiLU(), nn.LayerNorm(128), nn.Linear(128, 1536))
+        # self.last_linear2 = nn.Sequential(nn.Linear(667, 32))
+        self.skip_connection2 = nn.Linear(4096, 1, bias=False)
+        self.skip_connection = nn.Linear(667, 32, bias=False)
+        self.trans_linear = nn.Linear(666+1+4096, 1536, bias=False)
+        nn.init.constant_(self.skip_connection.weight.data, 0)
+        nn.init.constant_(self.trans_linear.weight.data, 0)
+        nn.init.constant_(self.trans_linear.weight.data, 0)
+        nn.init.constant_(self.pre8_linear[-1].weight.data, 0)
+        nn.init.constant_(self.pre16_linear[-1].weight.data, 0)
+        nn.init.constant_(self.pre24_linear[-1].weight.data, 0)
+        nn.init.constant_(self.pre8_linear2[-1].weight.data, 0)
+        nn.init.constant_(self.pre16_linear2[-1].weight.data, 0)
+        nn.init.constant_(self.pre24_linear2[-1].weight.data, 0)
+
+    def forward(self, noise: torch.Tensor, _s, _v, _d, _pool_embedding) -> torch.Tensor:
+        
+        assert noise is not None
+        _ori_v = _v.clone()
+        _v = torch.stack([torch.diag(_v[jj]) for jj in range(_v.shape[0])], dim=0)
+        positive_embedding = _s.permute(0, 2, 1) @ _v @ _d # [2, 64, 666] [2, 64] [2, 64, 4096]
+        pool_embedding = _pool_embedding[:, None, :]
+        embedding = torch.cat([positive_embedding, pool_embedding], dim=1)
+        bs = noise.shape[0]
+        height, width = noise.shape[-2:]
+        embed_8 = embedding
+        embed_16 = embedding
+        embed_24 = embedding
+        scale_8 = self.pre8_linear2(embed_8).mean(1)
+        scale_16 = self.pre16_linear2(embed_16).mean(1)
+        scale_24 = self.pre24_linear2(embed_24).mean(1)
+        embed_8 = self.pre8_linear(embed_8).mean(1)
+        embed_16 = self.pre16_linear(embed_16).mean(1)
+        embed_24 = self.pre24_linear(embed_24).mean(1)
+        embed_last = self.last_linear(embedding).mean(1)
+        embed_trans = self.trans_linear(torch.cat([_s, _ori_v[...,None], _d], dim=2)).mean(1)
+        skip_embedding = self.skip_connection(self.skip_connection2(embedding).permute(0,2,1)).permute(0,2,1)
+        scale_skip, embed_skip = skip_embedding.chunk(2,dim=1)
+        
+        ori_noise = noise * (scale_skip[...,None]) + embed_skip[...,None]
+        noise = self.pos_embed(noise)        
+        noise = noise * (1 + scale_8[:, None, :] + embed_trans[:, None, :]) + embed_8[:, None, :]
+        scale_list = [scale_16, scale_24]
+        embed_list = [embed_16, embed_24]
+        for _ii, block in enumerate(self.transformer_blocks):
+            noise = block(noise)  
+            if len(scale_list)!=0 and len(embed_list)!=0:
+                noise = noise * (1 + scale_list[int(_ii//4)][:, None, :] + embed_trans[:, None, :]) + embed_list[int(_ii//4)][:, None, :]
+        
+        hidden_states = noise
+        hidden_states = self.norm_out(hidden_states, embed_last)
+        hidden_states = self.proj_out(hidden_states)
+
+        # unpatchify
+        patch_size = self.patch_size
+        height = height // patch_size
+        width = width // patch_size
+    
+        hidden_states = hidden_states.reshape(
+            shape=(hidden_states.shape[0], height, width, patch_size, patch_size, self.out_channels)
+        )
+        hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
+        output = hidden_states.reshape(
+            shape=(hidden_states.shape[0], self.out_channels, height * patch_size, width * patch_size)
+        )
+        return output + ori_noise
+
+    def weak_load_state_dict(self, state_dict: os.Mapping[str, torch.any], strict: bool = True, assign: bool = False):
+        return load_filtered_state_dict(self, state_dict)
+
+class PromptSDXLNet(nn.Module):
+
+    def __init__(self,         
+        sample_size: int = 128,
+        patch_size: int = 2,
+        in_channels: int = 4,
+        num_layers: int = 4,
+        attention_head_dim: int = 64,
+        num_attention_heads: int = 24,
+        out_channels: int = 4,
+        pos_embed_max_size: int = 192
+        ):
+        super().__init__()
+        self.sample_size = sample_size
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.num_layers = num_layers
+        self.attention_head_dim = attention_head_dim
+        self.num_attention_heads = num_attention_heads
+        self.out_channels = out_channels
+        self.pos_embed_max_size = pos_embed_max_size
+        self.inner_dim = self.num_attention_heads * self.attention_head_dim
+        
+        self.pos_embed = PatchEmbed(
+            height=self.sample_size,
+            width=self.sample_size,
+            patch_size=self.patch_size,
+            in_channels=self.in_channels,
+            embed_dim=self.inner_dim,
+            pos_embed_max_size=pos_embed_max_size
+        )
+        
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.num_attention_heads,
+                    attention_head_dim=self.attention_head_dim,
+                    ff_inner_dim=2*self.inner_dim   # mult should be 4 by default
+                )
+                for i in range(self.num_layers)
+            ]
+        )
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+        
+        self.noise_shape = (1, 4, 128, 128) 
+        self.pre8_linear = nn.Sequential(nn.Linear(2048, 128), nn.SiLU(), nn.LayerNorm(128), nn.Linear(128, 1536))
+        self.pre16_linear = nn.Sequential(nn.Linear(2048, 128), nn.SiLU(), nn.LayerNorm(128), nn.Linear(128, 1536))
+        self.pre24_linear = nn.Sequential(nn.Linear(2048, 128), nn.SiLU(), nn.LayerNorm(128), nn.Linear(128, 1536))
+
+        self.pre8_linear2 = nn.Sequential(nn.Linear(2048, 128), nn.SiLU(), nn.LayerNorm(128), nn.Linear(128, 1536))
+        self.pre16_linear2 = nn.Sequential(nn.Linear(2048, 128), nn.SiLU(), nn.LayerNorm(128), nn.Linear(128, 1536))
+        self.pre24_linear2 = nn.Sequential(nn.Linear(2048, 128), nn.SiLU(), nn.LayerNorm(128), nn.Linear(128, 1536))
+        
+        self.last_linear = nn.Sequential(nn.Linear(2048, 128), nn.SiLU(), nn.LayerNorm(128), nn.Linear(128, 1536))
+        # self.last_linear2 = nn.Sequential(nn.Linear(667, 32))
+        self.skip_connection2 = nn.Linear(2048, 1, bias=False)
+        self.skip_connection = nn.Linear(154+1, 8, bias=False)
+        self.trans_linear = nn.Linear(154+1+2048, 1536, bias=False)
+        self.pool_prompt_linear = nn.Linear(2560, 2048, bias=False)
+        nn.init.constant_(self.skip_connection.weight.data, 0)
+        nn.init.constant_(self.trans_linear.weight.data, 0)
+        nn.init.constant_(self.trans_linear.weight.data, 0)
+        nn.init.constant_(self.pre8_linear[-1].weight.data, 0)
+        nn.init.constant_(self.pre16_linear[-1].weight.data, 0)
+        nn.init.constant_(self.pre24_linear[-1].weight.data, 0)
+        nn.init.constant_(self.pre8_linear2[-1].weight.data, 0)
+        nn.init.constant_(self.pre16_linear2[-1].weight.data, 0)
+        nn.init.constant_(self.pre24_linear2[-1].weight.data, 0)
+
+    def forward(self, noise: torch.Tensor, _s, _v, _d, _pool_embedding) -> torch.Tensor:
+        
+        assert noise is not None
+        _ori_v = _v.clone()
+        _v = torch.stack([torch.diag(_v[jj]) for jj in range(_v.shape[0])], dim=0)
+        positive_embedding = _s.permute(0, 2, 1) @ _v @ _d # [2, 64, 154] [2, 64] [2, 64, 2048]
+        pool_embedding = self.pool_prompt_linear(_pool_embedding[:, None, :])
+        embedding = torch.cat([positive_embedding, pool_embedding], dim=1)
+        bs = noise.shape[0]
+        height, width = noise.shape[-2:]
+        embed_8 = embedding
+        embed_16 = embedding
+        embed_24 = embedding
+        scale_8 = self.pre8_linear2(embed_8).mean(1)
+        scale_16 = self.pre16_linear2(embed_16).mean(1)
+        scale_24 = self.pre24_linear2(embed_24).mean(1)
+        embed_8 = self.pre8_linear(embed_8).mean(1)
+        embed_16 = self.pre16_linear(embed_16).mean(1)
+        embed_24 = self.pre24_linear(embed_24).mean(1)
+        embed_last = self.last_linear(embedding).mean(1)
+        embed_trans = self.trans_linear(torch.cat([_s, _ori_v[...,None], _d], dim=2)).mean(1)
+        skip_embedding = self.skip_connection(self.skip_connection2(embedding).permute(0,2,1)).permute(0,2,1)
+        scale_skip, embed_skip = skip_embedding.chunk(2,dim=1)
+        
+        ori_noise = noise * (scale_skip[...,None]) + embed_skip[...,None]
+        noise = self.pos_embed(noise)        
+        noise = noise * (1 + scale_8[:, None, :] + embed_trans[:, None, :]) + embed_8[:, None, :]
+        scale_list = [scale_16, scale_24]
+        embed_list = [embed_16, embed_24]
+        for _ii, block in enumerate(self.transformer_blocks):
+            noise = block(noise)  
+            if len(scale_list)!=0 and len(embed_list)!=0:
+                noise = noise * (1 + scale_list[int(_ii//4)][:, None, :] + embed_trans[:, None, :]) + embed_list[int(_ii//4)][:, None, :]
+        
+        hidden_states = noise
+        hidden_states = self.norm_out(hidden_states, embed_last)
+        hidden_states = self.proj_out(hidden_states)
+
+        # unpatchify
+        patch_size = self.patch_size
+        height = height // patch_size
+        width = width // patch_size
+    
+        hidden_states = hidden_states.reshape(
+            shape=(hidden_states.shape[0], height, width, patch_size, patch_size, self.out_channels)
+        )
+        hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
+        output = hidden_states.reshape(
+            shape=(hidden_states.shape[0], self.out_channels, height * patch_size, width * patch_size)
+        )
+        return output + ori_noise
+
+    def weak_load_state_dict(self, state_dict: os.Mapping[str, torch.any], strict: bool = True, assign: bool = False):
+        return load_filtered_state_dict(self, state_dict)
+
+
+class NoisePromptDataset(Dataset):
+    def __init__(self, if_weight=False):
+        
+        self.if_weight = if_weight
+        json_list = glob('/home/xiedian/total_datacollect/json/*.json')
+        self.original_score = []
+        self.optim_score = []
+        self.prompt = []
+        self.noise_paths = []
+        self.mask_conditions = []
+        self.embeddings = []
+        counter = 0
+        for i in range(len(json_list)):
+            with open('//home/xiedian/total_datacollect/json/new{:06d}.json'.format(i), 'r') as f:
+                data = json.load(f)
+            self.original_score.append(data['original_score_list'])
+            self.optim_score.append(data['optimized_score_list'])
+            if data['optimized_score_list']>data['original_score_list']:
+                counter += 1
+            self.prompt.append(data['caption'])
+            self.noise_paths.append('/home/xiedian/total_datacollect/latents/{:06d}.pt'.format(i))
+            self.embeddings.append('/data/xiedian/hg_reward/newdatacollect/total_datacollect_step_1_1/embedding/embeds_{:06d}.pt'.format(i))
+            z = [0, 1] * ((512+77+77) // 2)
+            self.mask_conditions.append(data['mid_token_ids'] if 'mid_token_ids' in data else z)
+        # while counter * 2 > len(self.prompt):
+        #     p = random.randint(0,len(self.prompt)-1)
+        #     if self.original_score[p] >  self.optim_score[p]:
+        #         self.optim_score.append(self.optim_score[p])
+        #         self.original_score.append(self.original_score[p])
+        #         self.mask_conditions.append(self.mask_conditions[p])
+        #         self.noise_paths.append(self.noise_paths[p])
+        #         self.prompt.append(self.prompt[p])
+        
+        # while counter * 2 < len(self.prompt):
+        #     p = random.randint(0,len(self.prompt)-1)
+        #     if self.original_score[p] >  self.optim_score[p]:
+        #         self.optim_score.append(self.optim_score[p])
+        #         self.original_score.append(self.original_score[p])
+        #         self.mask_conditions.append(self.mask_conditions[p])
+        #         self.noise_paths.append(self.noise_paths[p])
+        #         self.prompt.append(self.prompt[p])
+        
+        self.original_score = torch.Tensor(self.original_score)
+        self.optim_score = torch.Tensor(self.optim_score)
+                
+    def __len__(self):
+        return len(self.prompt)
+
+    def __getitem__(self, index):
+        try:
+            noise = torch.load(self.noise_paths[index], map_location='cpu').squeeze(0).float()
+            noise_pred_uncond, mid_noise_pred, noise_pred_text = noise.chunk(3,dim=0)
+            prompt = self.prompt[index]
+            original_score = self.original_score[index]
+            optim_score = self.optim_score[index]
+            embedding = torch.load(self.embeddings[index], map_location='cpu')
+            _s, _v, _d, _pool_embedding = embedding['_s'], embedding['_v'], embedding['_d'], embedding['pooled_prompt_embeds']
+            _s = _s.detach().float()
+            _v = _v.detach().float()
+            _d = _d.detach().float()
+            _pool_embedding = _pool_embedding.detach().float()
+            if original_score > optim_score:
+                noise_pred = noise_pred_uncond + 4.5 * (noise_pred_text - noise_pred_uncond)
+            else:
+                guidance_scale = 4.5 * 1.6
+                diff_text = torch.norm(noise_pred_text - noise_pred_uncond)
+                mid_guidance_scale = (diff_text / torch.norm(noise_pred_text - mid_noise_pred)).item()
+                guidance_scale_mid = guidance_scale / (2.4 + 1)
+                guidance_scale_all = guidance_scale * 2.4 / (2.4 + 1)
+                all_mid = (noise_pred_text - mid_noise_pred) * mid_guidance_scale
+                all_null = noise_pred_text - noise_pred_uncond
+                noise_pred = all_mid * guidance_scale_mid + all_null * guidance_scale_all + (mid_noise_pred + noise_pred_uncond) / 2
+        except:
+            print("error", index)
+            return self.__getitem__((index+1)%len(self.prompt))
+        if self.if_weight:
+            return noise_pred_text, prompt, noise_pred, 2 / (1+ math.exp((-abs(original_score) + 26.5)/1.7)), _s, _v, _d, _pool_embedding
+        else:
+            return noise_pred_text, prompt, noise_pred, _s, _v, _d, _pool_embedding
+
+class NoisePromptDataset_2_0(Dataset):
+    def __init__(self, if_weight=False):
+        
+        self.if_weight = if_weight
+        json_list = glob('/data/xiedian/hg_reward/newdatacollect/total_datacollect_step_1_1/json/*.json')
+        self.original_score = []
+        self.quick_score = []
+        self.slow_score = []
+        self.prompt = []
+        self.noise_paths = []
+        self.mask_conditions = []
+        self.img_list = []
+        self.embeddings = []
+        
+        counter = 0
+        for i in range(len(json_list)):
+            with open('/data/xiedian/hg_reward/newdatacollect/total_datacollect_step_1_1/json/new{:06d}.json'.format(i), 'r') as f:
+                data = json.load(f)
+            if (not os.path.exists('/data/xiedian/hg_reward/newdatacollect/total_datacollect_step_1_1/latents/{:06d}.pt'.format(i))) or \
+                max(data['original_score_list'], data['quick_score_list'], data['slow_score_list']) != data['original_score_list']:
+                continue
+            self.original_score.append(data['original_score_list'])
+            self.quick_score.append(data['quick_score_list'])
+            self.slow_score.append(data['slow_score_list'])
+            self.prompt.append(data['caption'])
+            self.noise_paths.append('/data/xiedian/hg_reward/newdatacollect/total_datacollect_step_1_1/latents/{:06d}.pt'.format(i))
+            z = [0, 1] * ((512+77+77) // 2)
+            self.mask_conditions.append(data['mid_token_ids'] if 'mid_token_ids' in data else z)
+            if data['original_score_list'] >= max(data['quick_score_list'], data['slow_score_list']):
+                self.img_list.append('/data/xiedian/hg_reward/newdatacollect/total_datacollect_step_1_1/optim/original{:06d}.png'.format(i))
+            elif data['quick_score_list'] >= max(data['original_score_list'], data['slow_score_list']):
+                self.img_list.append('/data/xiedian/hg_reward/newdatacollect/total_datacollect_step_1_1/optim/quick{:06d}.png'.format(i))
+            else:
+                self.img_list.append('/data/xiedian/hg_reward/newdatacollect/total_datacollect_step_1_1/optim/slow{:06d}.png'.format(i))
+            self.embeddings.append('/home/xiedian/total_datacollect/embedding/embeds_{:06d}.pt'.format(i))
+        self.original_score = torch.Tensor(self.original_score)
+        self.quick_score = torch.Tensor(self.quick_score)
+        self.slow_score = torch.Tensor(self.slow_score)
+            
+    def __len__(self):
+        return len(self.prompt)
+
+    def __getitem__(self, index):
+        try:
+            original_score = self.original_score[index]
+            quick_score = self.quick_score[index]
+            slow_score = self.slow_score[index]
+            original_score = max(max(quick_score, slow_score), original_score)
+            embedding = torch.load(self.embeddings[index], map_location='cpu')
+            _s, _v, _d, _pool_embedding = embedding['_s'], embedding['_v'], embedding['_d'], embedding['pooled_prompt_embeds']
+            _s = _s.detach().float()
+            _v = _v.detach().float()
+            _d = _d.detach().float()
+            _pool_embedding = _pool_embedding.detach().float()
+            noise = torch.load(self.noise_paths[index], map_location='cpu').squeeze(0).float()
+            noise_pred_text, noise_pred = noise.chunk(2,dim=0)
+            prompt = self.prompt[index]
+        except:
+            print("error", index)
+            return self.__getitem__((index+1)%len(self.prompt))
+        if self.if_weight:
+            return noise_pred_text, prompt, noise_pred, 2 / (1+ math.exp((-abs(original_score) + 26.5)/1.7)), _s, _v, _d, _pool_embedding
+        else:
+            return noise_pred_text, prompt, noise_pred, _s, _v, _d, _pool_embedding
+
+class NoisePromptDataset_3_0(Dataset):
+    def __init__(self, if_weight=False):
+        
+        self.if_weight = if_weight
+        json_list = glob('/data/xiedian/hg_reward/CFG_TOTAL/total_datacollect/json/*.json')
+        self.score = []
+        self.prompt = []
+        self.noise_paths = []
+        self.mask_conditions = []
+        self.img_list = []
+        self.embeddings = []
+        
+        print(len(json_list))
+        
+        for i in range(len(json_list)):
+            if (not os.path.exists('/data/xiedian/hg_reward/CFG_TOTAL/total_datacollect/embedding/{:06d}.pt'.format(i))):
+                continue
+            
+            with open('/data/xiedian/hg_reward/CFG_TOTAL/total_datacollect/json/new{:06d}.json'.format(i), 'r') as f:
+                data = json.load(f)
+            if data['original_score_list'] > data['optimized_score_list']:
+                tag = 0
+                if (not os.path.exists('/data/xiedian/hg_reward/CFG_TOTAL/total_datacollect/latents/original{:06d}.pt'.format(i))):
+                    continue
+            else:
+                tag = 1
+                if (not os.path.exists('/data/xiedian/hg_reward/CFG_TOTAL/total_datacollect/latents/new{:06d}.pt'.format(i))):
+                    continue
+            
+            if tag == 1:
+                self.score.append(data['optimized_score_list'])
+                self.noise_paths.append('/data/xiedian/hg_reward/CFG_TOTAL/total_datacollect/latents/new{:06d}.pt'.format(i))
+            else:
+                self.score.append(data['original_score_list'])
+                self.noise_paths.append('/data/xiedian/hg_reward/CFG_TOTAL/total_datacollect/latents/original{:06d}.pt'.format(i))
+            self.prompt.append(data['caption'])
+            self.embeddings.append('/data/xiedian/hg_reward/CFG_TOTAL/total_datacollect/embedding/{:06d}.pt'.format(i))
+        self.score = torch.Tensor(self.score)
+            
+    def __len__(self):
+        return len(self.prompt)
+
+    def __getitem__(self, index):
+        try:
+            embedding = torch.load(self.embeddings[index], map_location='cpu')
+            _s, _v, _d, _pool_embedding = embedding['_s'], embedding['_v'], embedding['_d'], embedding['_pooled_prompt_embeds']
+            _s = _s.detach().float()
+            _v = _v.detach().float()
+            _d = _d.detach().float()
+            _pool_embedding = _pool_embedding.detach().float()
+            noise = torch.load(self.noise_paths[index], map_location='cpu').float() # [2XT, 16, 128, 128]
+            prompt = self.prompt[index] # [ori, target, ori]
+            score = self.score[index]
+        except:
+            print("error", index)
+            return self.__getitem__((index+1)%len(self.prompt))
+        if self.if_weight:
+            return noise, prompt, 2 / (1+ math.exp((-abs(score) + 26.5)/1.7)), _s, _v, _d, _pool_embedding
+        else:
+            return noise, prompt, _s, _v, _d, _pool_embedding
+
+
+def load_filtered_state_dict(model, state_dict):
+    model_state_dict = model.state_dict()
+    filtered_state_dict = {}
+    for k, v in state_dict.items():
+        if k in model_state_dict:
+            if model_state_dict[k].size() == v.size():
+                filtered_state_dict[k] = v
+            else:
+                print(f"Skipping {k}: shape mismatch ({model_state_dict[k].size()} vs {v.size()})")
+        else:
+            print(f"Skipping {k}: not found in model's state_dict.")
+    model.load_state_dict(filtered_state_dict, strict=False)
+    return model
+
+def custom_collate_fn_2_0(batch):
+    noise_pred_texts, prompts, noise_preds, max_scores = zip(*batch)
+    
+    noise_pred_texts = torch.stack(noise_pred_texts)
+    noise_preds = torch.stack(noise_preds)
+    max_scores = torch.stack(max_scores)
+    
+    return noise_pred_texts, prompts, noise_preds, max_scores
+
+
+if __name__ == "__main__":
+    dataset = NoisePromptDataset(if_weight=True)
+    weights = []
+    for i, (noise, prompt, gt, weight) in enumerate(dataset):
+        weights.append(weight)
+    weights = torch.from_numpy(np.array(weights)).cuda()
+    print(weights.mean(), weights.std(dim=0))

sample_img.py

@@ -0,0 +1,237 @@
+import json
+import numpy as np
+import math
+import csv
+import random
+import argparse
+import torch
+import os
+import torch.distributed as dist
+
+from PIL import Image
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from accelerate.utils import set_seed
+
+from diffusion_pipeline.sd35_pipeline import StableDiffusion3Pipeline, FlowMatchEulerInverseScheduler
+from diffusion_pipeline.sdxl_pipeline import StableDiffusionXLPipeline
+from diffusers import BitsAndBytesConfig, SD3Transformer2DModel
+from diffusers import FlowMatchEulerDiscreteScheduler, DDIMInverseScheduler, DDIMScheduler
+
+device = torch.device('cuda')
+
+def get_args():
+      # pick: test_unique_caption_zh.csv       draw: drawbench.csv
+      parser = argparse.ArgumentParser()
+      parser.add_argument("--model", default='sd35', choices=['sdxl', 'sd35'], type=str)
+      parser.add_argument("--inference-step", default=30, type=int)
+      parser.add_argument("--size", default=1024, type=int)
+      parser.add_argument("--seed", default=33, type=int)
+      parser.add_argument("--cfg", default=3.5, type=float)
+
+      # hyperparameters for Z-Sampling
+      parser.add_argument("--inv-cfg", default=0.5, type=float)
+
+      # hyperparameters for Z-Core^2
+      parser.add_argument("--w2s-guidance", default=1.5, type=float)
+      parser.add_argument("--end_timesteps", default=28, type=int) # equal to inference step - 2 or inference step
+
+
+      parser.add_argument("--prompt", default='Mickey Mouse painting by Frank Frazetta.', type=str)
+
+      parser.add_argument("--method", default='standard', choices=['standard', 'core', 'zigzag', 'z-core'], type=str)
+
+      args =  parser.parse_args()
+      return args
+
+
+if __name__ == '__main__':
+      torch.cuda.empty_cache()
+      dtype = torch.float16
+      args = get_args()
+      print("args.seed: ", args.seed)
+      set_seed(args.seed)
+
+      # TODO: load pipeline
+      if args.model == 'sd35':
+            nf4_config = BitsAndBytesConfig(
+                  load_in_4bit=True,
+                  bnb_4bit_quant_type="nf4",
+                  bnb_4bit_compute_dtype=torch.bfloat16
+            )
+            model_nf4 = SD3Transformer2DModel.from_pretrained(
+                  "stabilityai/stable-diffusion-3.5-large",
+                  subfolder="transformer",
+                  quantization_config=nf4_config,
+                  torch_dtype=torch.bfloat16
+            )
+
+            pipe = StableDiffusion3Pipeline.from_pretrained(
+                  "stabilityai/stable-diffusion-3.5-large", 
+                  transformer=model_nf4,
+                  torch_dtype=torch.bfloat16,
+            )
+
+            pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config)
+            inverse_scheduler = FlowMatchEulerInverseScheduler.from_pretrained("stabilityai/stable-diffusion-3.5-large",
+                                                                  subfolder='scheduler')
+            pipe.inv_scheduler = inverse_scheduler
+
+      elif args.model == "sdxl":
+            pipe = StableDiffusionXLPipeline.from_pretrained(
+                  "stabilityai/stable-diffusion-xl-base-1.0", 
+                  torch_dtype=torch.float16, 
+                  variant="fp16",
+                  use_safetensors=True
+            ).to("cuda")
+
+            pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+            inverse_scheduler = DDIMInverseScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0",
+                                                                  subfolder='scheduler')
+            pipe.inv_scheduler = inverse_scheduler
+
+      pipe.to(device)
+      pipe.enable_model_cpu_offload()
+
+      # TODO: load noise model
+      if args.method == 'core' or args.method == 'z-core':
+            from diffusion_pipeline.refine_model import PromptSD35Net, PromptSDXLNet
+            from diffusion_pipeline.lora import replace_linear_with_lora, lora_true
+
+            if args.model == 'sd35':
+                  refine_model = PromptSD35Net()
+                  replace_linear_with_lora(refine_model, rank=64, alpha=1.0, number_of_lora=28)
+                  lora_true(refine_model, lora_idx=0)
+
+                  checkpoint = torch.load('./weights/sd35_ckpt_v9.pth', map_location='cpu')
+                  refine_model.load_state_dict(checkpoint)
+            elif args.model == 'sdxl':
+                  refine_model = PromptSDXLNet()
+                  replace_linear_with_lora(refine_model, rank=48, alpha=1.0, number_of_lora=50)
+                  lora_true(refine_model, lora_idx=0)
+
+                  checkpoint = torch.load('./weights/sdxl_ckpt_v9.pth', map_location='cpu')
+                  refine_model.load_state_dict(checkpoint)
+
+            print("Load Lora Success")
+            refine_model = refine_model.to(device)
+            refine_model = refine_model.to(torch.bfloat16)
+
+      
+      # TODO: load hyperparameters
+      size = args.size
+      if args.model == 'sdxl':
+            shape = (1, 4, size // 8, size // 8)
+      else:
+            shape = (1, 16, size // 8, size // 8)
+
+      num_steps = args.inference_step
+      end_timesteps = args.end_timesteps
+      guidance_scale = args.cfg
+      w2s_guidance = args.w2s_guidance
+      inv_cfg = args.inv_cfg
+      prompt = args.prompt
+
+      print("pass this prompt: ", prompt)
+    
+      start_latents = torch.randn(shape, dtype=dtype).to(device)
+        
+      if args.model == 'sdxl':
+            if args.method == 'core':
+                  output = pipe.core(
+                        prompt=prompt,
+                        guidance_scale=guidance_scale,
+                        num_inference_steps=num_steps,
+                        latents=start_latents,
+                        return_dict=False,
+                        refine_model=refine_model,
+                        lora_true=lora_true,
+                        end_timesteps=end_timesteps,
+                        w2s_guidance=w2s_guidance)[0][0]  
+            
+            elif args.method == 'zigzag':
+                  output = pipe.zigzag(
+                        prompt=prompt,
+                        guidance_scale=guidance_scale,
+                        latents=start_latents,
+                        return_dict=False,
+                        num_inference_steps=num_steps,
+                        inv_cfg=inv_cfg)[0][0] 
+
+            elif args.method == 'z-core':
+                  output = pipe.z_core(
+                        prompt=prompt,
+                        guidance_scale=guidance_scale,
+                        num_inference_steps=num_steps,
+                        latents=start_latents,
+                        return_dict=False,
+                        refine_model=refine_model,
+                        lora_true=lora_true,
+                        end_timesteps=end_timesteps,
+                        w2s_guidance=w2s_guidance,
+                        inv_cfg=inv_cfg)[0][0]  
+                  
+            elif args.method == 'standard':
+                  output = pipe(
+                        prompt=prompt,
+                        guidance_scale=guidance_scale,
+                        latents=start_latents,
+                        return_dict=False,
+                        num_inference_steps=num_steps)[0][0] 
+            else:
+                  raise ValueError("Invalid method")
+            
+            output.save(f'{args.model}_{args.method}.png')  
+
+
+      else: 
+            if args.method == 'core':
+                  output = pipe.core(
+                        prompt=prompt,
+                        guidance_scale=guidance_scale,
+                        num_inference_steps=num_steps,
+                        latents=start_latents,
+                        max_sequence_length=512,
+                        return_dict=False,
+                        refine_model=refine_model,
+                        lora_true=lora_true,
+                        end_timesteps=end_timesteps,
+                        w2s_guidance=w2s_guidance)[0][0]  
+            
+            elif args.method == 'zigzag':
+                  output = pipe.zigzag(
+                        prompt=prompt,
+                        max_sequence_length=512,
+                        guidance_scale=guidance_scale,
+                        latents=start_latents,
+                        return_dict=False,
+                        num_inference_steps=num_steps,
+                        inv_cfg=inv_cfg)[0][0] 
+
+            elif args.method == 'z-core':
+                  output = pipe.z_core(
+                        prompt=prompt,
+                        guidance_scale=guidance_scale,
+                        num_inference_steps=num_steps,
+                        latents=start_latents,
+                        return_dict=False,
+                        max_sequence_length=512,
+                        refine_model=refine_model,
+                        lora_true=lora_true,
+                        end_timesteps=end_timesteps,
+                        w2s_guidance=w2s_guidance)[0][0]   
+                  
+            elif args.method == 'standard':
+                  output = pipe(
+                        prompt=prompt,
+                        guidance_scale=guidance_scale,
+                        latents=start_latents,
+                        return_dict=False,
+                        max_sequence_length=512,
+                        num_inference_steps=num_steps)[0][0] 
+            else:
+                  raise ValueError("Invalid method")
+            
+            output.save(f'{args.model}_{args.method}.png')  
+
+