| import torch | |
| import numpy as np | |
| from .processors import Processor_id | |
| class ControlNetConfigUnit: | |
| def __init__(self, processor_id: Processor_id, model_path, scale=1.0, skip_processor=False): | |
| self.processor_id = processor_id | |
| self.model_path = model_path | |
| self.scale = scale | |
| self.skip_processor = skip_processor | |
| class ControlNetUnit: | |
| def __init__(self, processor, model, scale=1.0): | |
| self.processor = processor | |
| self.model = model | |
| self.scale = scale | |
| class MultiControlNetManager: | |
| def __init__(self, controlnet_units=[]): | |
| self.processors = [unit.processor for unit in controlnet_units] | |
| self.models = [unit.model for unit in controlnet_units] | |
| self.scales = [unit.scale for unit in controlnet_units] | |
| def cpu(self): | |
| for model in self.models: | |
| model.cpu() | |
| def to(self, device): | |
| for model in self.models: | |
| model.to(device) | |
| for processor in self.processors: | |
| processor.to(device) | |
| def process_image(self, image, processor_id=None): | |
| if processor_id is None: | |
| processed_image = [processor(image) for processor in self.processors] | |
| else: | |
| processed_image = [self.processors[processor_id](image)] | |
| processed_image = torch.concat([ | |
| torch.Tensor(np.array(image_, dtype=np.float32) / 255).permute(2, 0, 1).unsqueeze(0) | |
| for image_ in processed_image | |
| ], dim=0) | |
| return processed_image | |
| def __call__( | |
| self, | |
| sample, timestep, encoder_hidden_states, conditionings, | |
| tiled=False, tile_size=64, tile_stride=32, **kwargs | |
| ): | |
| res_stack = None | |
| for processor, conditioning, model, scale in zip(self.processors, conditionings, self.models, self.scales): | |
| res_stack_ = model( | |
| sample, timestep, encoder_hidden_states, conditioning, **kwargs, | |
| tiled=tiled, tile_size=tile_size, tile_stride=tile_stride, | |
| processor_id=processor.processor_id | |
| ) | |
| res_stack_ = [res * scale for res in res_stack_] | |
| if res_stack is None: | |
| res_stack = res_stack_ | |
| else: | |
| res_stack = [i + j for i, j in zip(res_stack, res_stack_)] | |
| return res_stack | |
| class FluxMultiControlNetManager(MultiControlNetManager): | |
| def __init__(self, controlnet_units=[]): | |
| super().__init__(controlnet_units=controlnet_units) | |
| def process_image(self, image, processor_id=None): | |
| if processor_id is None: | |
| processed_image = [processor(image) for processor in self.processors] | |
| else: | |
| processed_image = [self.processors[processor_id](image)] | |
| return processed_image | |
| def __call__(self, conditionings, **kwargs): | |
| res_stack, single_res_stack = None, None | |
| for processor, conditioning, model, scale in zip(self.processors, conditionings, self.models, self.scales): | |
| res_stack_, single_res_stack_ = model(controlnet_conditioning=conditioning, processor_id=processor.processor_id, **kwargs) | |
| res_stack_ = [res * scale for res in res_stack_] | |
| single_res_stack_ = [res * scale for res in single_res_stack_] | |
| if res_stack is None: | |
| res_stack = res_stack_ | |
| single_res_stack = single_res_stack_ | |
| else: | |
| res_stack = [i + j for i, j in zip(res_stack, res_stack_)] | |
| single_res_stack = [i + j for i, j in zip(single_res_stack, single_res_stack_)] | |
| return res_stack, single_res_stack | |