98 lines
4.4 KiB
Python
98 lines
4.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import torch
|
|
from torch.nn.utils.rnn import pad_sequence
|
|
|
|
from einops import rearrange
|
|
from scepter.modules.model.backbone.flux import FluxMR
|
|
from scepter.modules.model.registry import BACKBONES
|
|
from scepter.modules.utils.config import dict_to_yaml
|
|
|
|
|
|
@BACKBONES.register_class()
|
|
class FluxMRACEPlus(FluxMR):
|
|
def __init__(self, cfg, logger=None):
|
|
super().__init__(cfg, logger)
|
|
|
|
def prepare_input(self, x, cond):
|
|
context, y = cond['context'], cond['y']
|
|
batch_frames, batch_frames_ids = [], []
|
|
for ix, shape, imask, ie, ie_mask in zip(x, cond['x_shapes'],
|
|
cond['x_mask'], cond['edit'],
|
|
cond['edit_mask']):
|
|
# unpack image from sequence
|
|
ix = ix[:, :shape[0] * shape[1]].view(-1, shape[0], shape[1])
|
|
imask = torch.ones_like(
|
|
ix[[0], :, :]) if imask is None else imask.squeeze(0)
|
|
if len(ie) > 0:
|
|
ie = [iie.squeeze(0) for iie in ie]
|
|
ie_mask = [
|
|
torch.ones(
|
|
(ix.shape[0] * 4, ix.shape[1],
|
|
ix.shape[2])) if iime is None else iime.squeeze(0)
|
|
for iime in ie_mask
|
|
]
|
|
ie = torch.cat(ie, dim=-1)
|
|
ie_mask = torch.cat(ie_mask, dim=-1)
|
|
else:
|
|
ie, ie_mask = torch.zeros_like(ix).to(x), torch.ones_like(
|
|
imask).to(x)
|
|
ix = torch.cat([ix, ie, ie_mask], dim=0)
|
|
c, h, w = ix.shape
|
|
ix = rearrange(ix,
|
|
'c (h ph) (w pw) -> (h w) (c ph pw)',
|
|
ph=2,
|
|
pw=2)
|
|
ix_id = torch.zeros(h // 2, w // 2, 3)
|
|
ix_id[..., 1] = ix_id[..., 1] + torch.arange(h // 2)[:, None]
|
|
ix_id[..., 2] = ix_id[..., 2] + torch.arange(w // 2)[None, :]
|
|
ix_id = rearrange(ix_id, 'h w c -> (h w) c')
|
|
batch_frames.append([ix])
|
|
batch_frames_ids.append([ix_id])
|
|
x_list, x_id_list, mask_x_list, x_seq_length = [], [], [], []
|
|
for frames, frame_ids in zip(batch_frames, batch_frames_ids):
|
|
proj_frames = []
|
|
for idx, one_frame in enumerate(frames):
|
|
one_frame = self.img_in(one_frame)
|
|
proj_frames.append(one_frame)
|
|
ix = torch.cat(proj_frames, dim=0)
|
|
if_id = torch.cat(frame_ids, dim=0)
|
|
x_list.append(ix)
|
|
x_id_list.append(if_id)
|
|
mask_x_list.append(
|
|
torch.ones(ix.shape[0]).to(ix.device,
|
|
non_blocking=True).bool())
|
|
x_seq_length.append(ix.shape[0])
|
|
# if len(x_list) < 1: import pdb;pdb.set_trace()
|
|
x = pad_sequence(tuple(x_list), batch_first=True)
|
|
x_ids = pad_sequence(tuple(x_id_list), batch_first=True).to(
|
|
x) # [b,pad_seq,2] pad (0.,0.) at dim2
|
|
mask_x = pad_sequence(tuple(mask_x_list), batch_first=True)
|
|
# import pdb;pdb.set_trace()
|
|
if isinstance(context, list):
|
|
txt_list, mask_txt_list, y_list = [], [], []
|
|
for sample_id, (ctx, yy) in enumerate(zip(context, y)):
|
|
txt_list.append(self.txt_in(ctx.to(x)))
|
|
mask_txt_list.append(
|
|
torch.ones(txt_list[-1].shape[0]).to(
|
|
ctx.device, non_blocking=True).bool())
|
|
y_list.append(yy.to(x))
|
|
txt = pad_sequence(tuple(txt_list), batch_first=True)
|
|
txt_ids = torch.zeros(txt.shape[0], txt.shape[1], 3).to(x)
|
|
mask_txt = pad_sequence(tuple(mask_txt_list), batch_first=True)
|
|
y = torch.cat(y_list, dim=0)
|
|
assert y.ndim == 2 and txt.ndim == 3
|
|
else:
|
|
txt = self.txt_in(context)
|
|
txt_ids = torch.zeros(context.shape[0], context.shape[1], 3).to(x)
|
|
mask_txt = torch.ones(context.shape[0], context.shape[1]).to(
|
|
x.device, non_blocking=True).bool()
|
|
return x, x_ids, txt, txt_ids, y, mask_x, mask_txt, x_seq_length
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
return dict_to_yaml('MODEL',
|
|
__class__.__name__,
|
|
FluxMRACEPlus.para_dict,
|
|
set_name=True)
|