Paint by Example 논문 리뷰, Paint by Example 써보기
마스크 이미지를 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)
'코드 리뷰 > Diffusion' 카테고리의 다른 글
CompVis/stable-diffusion/main.py (0) | 2023.01.19 |
---|---|
DiffStyler 코드 리뷰 (0) | 2023.01.16 |
DAAM 코드 리뷰 (0) | 2023.01.12 |
Latent Diffusion (0) | 2022.12.28 |
Classifier-Guidance Diffusion (1) | 2022.12.07 |
Improved DDPM (0) | 2022.09.28 |