| 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): | |
| self.processor_id = processor_id | |
| self.model_path = model_path | |
| self.scale = scale | |
| 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 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 | |
| ): | |
| res_stack = None | |
| for conditioning, model, scale in zip(conditionings, self.models, self.scales): | |
| res_stack_ = model( | |
| sample, timestep, encoder_hidden_states, conditioning, | |
| tiled=tiled, tile_size=tile_size, tile_stride=tile_stride | |
| ) | |
| 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 | |