81 lines
3.3 KiB
Python
81 lines
3.3 KiB
Python
from ..utils.module_utils import is_named_module_transformer_block
|
|
from .ref_mode import RefMode
|
|
from .ref_config import RefConfig
|
|
|
|
|
|
class RefController:
|
|
model = None
|
|
|
|
def __init__(self, diffusion_model):
|
|
self.model = diffusion_model
|
|
|
|
def sef_ref_count(self, ref_count):
|
|
model = self.model
|
|
input_modules = list(
|
|
filter(is_named_module_transformer_block, model.input_blocks.named_modules()))
|
|
output_modules = list(
|
|
filter(is_named_module_transformer_block, model.output_blocks.named_modules()))
|
|
|
|
for _, module in input_modules + output_modules:
|
|
for i in range(len(module.attention_blocks)):
|
|
attn = module.attention_blocks[i]
|
|
attn.ref_count = ref_count
|
|
|
|
def set_motion_mode(self, mode: RefMode, config: RefConfig = None):
|
|
model = self.model
|
|
input_modules = list(
|
|
filter(is_named_module_transformer_block, model.input_blocks.named_modules()))
|
|
output_modules = list(
|
|
filter(is_named_module_transformer_block, model.output_blocks.named_modules()))
|
|
|
|
for _, module in input_modules + output_modules:
|
|
for i in range(len(module.attention_blocks)):
|
|
attn = module.attention_blocks[i]
|
|
attn.ref_off()
|
|
|
|
if config is None or len(config.settings) == 0:
|
|
return
|
|
|
|
for setting in config.settings:
|
|
for i, (_, module) in enumerate(input_modules):
|
|
for attn_idx in range(len(module.attention_blocks)):
|
|
attn = module.attention_blocks[attn_idx]
|
|
if setting.input_attentions[i]:
|
|
attn.ref_mode = mode
|
|
if setting.q_mode:
|
|
attn.q_mode = True
|
|
if setting.k_mode:
|
|
attn.k_mode = True
|
|
if setting.v_mode:
|
|
attn.v_mode = True
|
|
if setting.normal_mode:
|
|
attn.normal_mode = True
|
|
attn.ref_norm_fidelity = setting.normal_fidelity
|
|
for i, (_, module) in enumerate(output_modules):
|
|
for attn_idx in range(len(module.attention_blocks)):
|
|
attn = module.attention_blocks[attn_idx]
|
|
if setting.output_attentions[i]:
|
|
attn.ref_mode = mode
|
|
if setting.q_mode:
|
|
attn.q_mode = True
|
|
if setting.k_mode:
|
|
attn.k_mode = True
|
|
if setting.v_mode:
|
|
attn.v_mode = True
|
|
if setting.normal_mode:
|
|
attn.normal_mode = True
|
|
attn.ref_norm_fidelity = setting.normal_fidelity
|
|
|
|
def clear_modules(self):
|
|
model = self.model
|
|
input_modules = list(
|
|
filter(is_named_module_transformer_block, model.input_blocks.named_modules()))
|
|
output_modules = list(
|
|
filter(is_named_module_transformer_block, model.output_blocks.named_modules()))
|
|
|
|
for _, module in input_modules + output_modules:
|
|
for i in range(len(module.attention_blocks)):
|
|
attn = module.attention_blocks[i]
|
|
attn.ref_off()
|
|
attn.ref_clean()
|