60 lines
1.7 KiB
Python
60 lines
1.7 KiB
Python
import enum
|
|
|
|
from .block_type import BlockType
|
|
|
|
|
|
class RefMode(enum.Enum):
|
|
OFF = 1
|
|
WRITE = 2
|
|
READ = 3
|
|
|
|
|
|
class RefController:
|
|
modules = []
|
|
temp_modules = []
|
|
default_video_length = 16
|
|
default_full_length = 16
|
|
|
|
def add_module(self, module):
|
|
self.modules.append(module)
|
|
|
|
def set_mode(self, mode: RefMode):
|
|
for module in self.modules:
|
|
module.ref_mode = mode
|
|
if mode == RefMode.OFF:
|
|
self.set_kv_mode(mode)
|
|
self.set_normal_mode(mode)
|
|
|
|
def set_normal_mode(self, mode: RefMode):
|
|
for module in self.modules:
|
|
module.normal_mode = mode
|
|
|
|
def set_kv_mode(self, mode: RefMode, count = None):
|
|
for module in self.modules:
|
|
module.kv_mode = RefMode.OFF
|
|
if mode == RefMode.OFF or count is None:
|
|
count = len(self.modules)
|
|
output_modules = list(filter(lambda m: m.block_type == BlockType.OUTPUT, self.modules))
|
|
output_modules = sorted(output_modules, key=lambda m: -m.block_idx)[:count]
|
|
for module in output_modules:
|
|
module.kv_mode = mode
|
|
|
|
def clear_modules(self):
|
|
for module in self.modules:
|
|
module.clean()
|
|
self.set_mode(RefMode.OFF)
|
|
self.set_kv_mode(RefMode.OFF, len(self.modules))
|
|
self.set_normal_mode(RefMode.OFF)
|
|
|
|
def set_cfg(self, cfg: bool):
|
|
for module in self.modules:
|
|
module.is_cfg = cfg
|
|
|
|
def set_attention_bound(self, attention_bound: float):
|
|
for module in self.modules:
|
|
module.attention_bound = attention_bound
|
|
|
|
def set_content_fidelity(self, content_fidelity: float):
|
|
for module in self.modules:
|
|
module.content_fidelity = content_fidelity
|