Paint by Example 코드 리뷰

Paint by Example 논문 리뷰, Paint by Example 써보기

---

Github

GitHub - Fantasy-Studio/Paint-by-Example: Paint by Example: Exemplar-based Image Editing with Diffusion Models

 

---

마스크 이미지를 U-Net에 전달하기 위해 기존 모델에 추가 채널을 받는 것은 간단하다.

import torch
pretrained_model_path='pretrained_models/sd-v1-4.ckpt'
ckpt_file=torch.load(pretrained_model_path,map_location='cpu')
zero_data=torch.zeros(320,5,3,3)
new_weight=torch.cat((ckpt_file['state_dict']['model.diffusion_model.input_blocks.0.0.weight'],zero_data),dim=1)
ckpt_file['state_dict']['model.diffusion_model.input_blocks.0.0.weight']=new_weight
torch.save(ckpt_file,"pretrained_models/sd-v1-4-modified-9channel.ckpt")

---

데이터셋

이미지에서 bbox를 추출하고 증강, 그리고 bbox를 기준으로 변형된 마스크를 적용한 이미지가 필요하다.

 

위 데이터 전처리의 일련의 과정은 open_images.py에 있다.

마스크 형태를 변형시키는 방법은 다음과 같다.

 

다음 그림이 있을 때

 

먼저 bbox와 extended bbox를 정의하고

extended_bbox=copy.copy(bbox)
left_freespace=bbox[0]-0
right_freespace=W-bbox[2]
up_freespace=bbox[1]-0
down_freespace=H-bbox[3]
extended_bbox[0]=bbox[0]-random.randint(0,int(0.4*left_freespace))
extended_bbox[1]=bbox[1]-random.randint(0,int(0.4*up_freespace))
extended_bbox[2]=bbox[2]+random.randint(0,int(0.4*right_freespace))
extended_bbox[3]=bbox[3]+random.randint(0,int(0.4*down_freespace))

 

각 노드를 구하고 bezier curve를 그림

if prob<self.arbitrary_mask_percent:
    mask_img = Image.new('RGB', (W, H), (255, 255, 255)) 
    bbox_mask=copy.copy(bbox)
    extended_bbox_mask=copy.copy(extended_bbox)
    top_nodes = np.asfortranarray([
                    [bbox_mask[0],(bbox_mask[0]+bbox_mask[2])/2 , bbox_mask[2]],
                    [bbox_mask[1], extended_bbox_mask[1], bbox_mask[1]],
                ])
    down_nodes = np.asfortranarray([
            [bbox_mask[2],(bbox_mask[0]+bbox_mask[2])/2 , bbox_mask[0]],
            [bbox_mask[3], extended_bbox_mask[3], bbox_mask[3]],
        ])
    left_nodes = np.asfortranarray([
            [bbox_mask[0],extended_bbox_mask[0] , bbox_mask[0]],
            [bbox_mask[3], (bbox_mask[1]+bbox_mask[3])/2, bbox_mask[1]],
        ])
    right_nodes = np.asfortranarray([
            [bbox_mask[2],extended_bbox_mask[2] , bbox_mask[2]],
            [bbox_mask[1], (bbox_mask[1]+bbox_mask[3])/2, bbox_mask[3]],
        ])
    top_curve = bezier.Curve(top_nodes,degree=2)
    right_curve = bezier.Curve(right_nodes,degree=2)
    down_curve = bezier.Curve(down_nodes,degree=2)
    left_curve = bezier.Curve(left_nodes,degree=2)
    curve_list=[top_curve,right_curve,down_curve,left_curve]

 

Bezier curve에서 일정한 간격으로 점을 샘플링한 뒤에 무작위 픽셀 오프셋을 추가

for curve in curve_list:
    x_list=[]
    y_list=[]
    for i in range(1,19):
        if (curve.evaluate(i*0.05)[0][0]) not in x_list and (curve.evaluate(i*0.05)[1][0] not in y_list):
            pt_list.append((curve.evaluate(i*0.05)[0][0]+random.randint(-random_width,random_width),curve.evaluate(i*0.05)[1][0]+random.randint(-random_width,random_width)))
            x_list.append(curve.evaluate(i*0.05)[0][0])
            y_list.append(curve.evaluate(i*0.05)[1][0])

 

연결

mask_img_draw=ImageDraw.Draw(mask_img)
mask_img_draw.polygon(pt_list,fill=(0,0,0))
mask_tensor=get_tensor(normalize=False, toTensor=True)(mask_img)[0].unsqueeze(0)

---

CLIP 인코더에도 학습 가능한 최종 계층 추가

class FrozenCLIPImageEmbedder(AbstractEncoder):
    """Uses the CLIP transformer encoder for text (from Hugging Face)"""
    def __init__(self, version="openai/clip-vit-large-patch14"):
        super().__init__()
        self.transformer = CLIPVisionModel.from_pretrained(version)
        self.final_ln = LayerNorm(1024)
        self.mapper = Transformer(
                1,
                1024,
                5,
                1,
            )

        self.freeze()

    def freeze(self):
        self.transformer = self.transformer.eval()
        for param in self.parameters():
            param.requires_grad = False
        for param in self.mapper.parameters():
            param.requires_grad = True
        for param in self.final_ln.parameters():
            param.requires_grad = True

 

MLP 블록도 추가

class LatentDiffusion(DDPM):
	def __init__(self, *args, **kwargs):
    
    	...
        self.proj_out=nn.Linear(1024, 768)
        ...
        
    def forward(self, x, c, *args, **kwargs):
    
    	...
        if self.cond_stage_trainable:
            c = self.get_learned_conditioning(c)
            c = self.proj_out(c)
        ...

---

개선된 Classifier-free Guidance

 

학습 가능한 벡터 v 추가, 무조건 모델의 경우에 v를 조건으로 받음

class LatentDiffusion(DDPM):
	def __init__(self, *args, **kwargs):
    
    	...
        self.learnable_vector = nn.Parameter(torch.randn((1,1,768)), requires_grad=True)
        ...
        
    def forward(self, x, c, *args, **kwargs):
    
    	...
        if self.u_cond_prop<self.u_cond_percent:
            self.opt.params=self.params_with_white
            return self.p_losses(x, self.learnable_vector.repeat(x.shape[0],1,1), t, *args, **kwargs)
        else:
            return self.p_losses(x, c, t, *args, **kwargs)

 

 

optimizer에 MLP와 CLIP 인코더 최종 계층, learnable vector 추가

def configure_optimizers(self):

	...
    if self.cond_stage_trainable:
        print(f"{self.__class__.__name__}: Also optimizing conditioner params!")
        params = params + list(self.cond_stage_model.final_ln.parameters())+list(self.cond_stage_model.mapper.parameters())+list(self.proj_out.parameters())
  	
    self.params = params
    self.params_with_white=params + list(self.learnable_vector)
    opt = torch.optim.AdamW(params, lr=lr)