| 19 | |
| 20 | |
| 21 | class MultiControlNetManager: |
| 22 | def __init__(self, controlnet_units=[]): |
| 23 | self.processors = [unit.processor for unit in controlnet_units] |
| 24 | self.models = [unit.model for unit in controlnet_units] |
| 25 | self.scales = [unit.scale for unit in controlnet_units] |
| 26 | |
| 27 | def cpu(self): |
| 28 | for model in self.models: |
| 29 | model.cpu() |
| 30 | |
| 31 | def to(self, device): |
| 32 | for model in self.models: |
| 33 | model.to(device) |
| 34 | for processor in self.processors: |
| 35 | processor.to(device) |
| 36 | |
| 37 | def process_image(self, image, processor_id=None): |
| 38 | if processor_id is None: |
| 39 | processed_image = [processor(image) for processor in self.processors] |
| 40 | else: |
| 41 | processed_image = [self.processors[processor_id](image)] |
| 42 | processed_image = torch.concat([ |
| 43 | torch.Tensor(np.array(image_, dtype=np.float32) / 255).permute(2, 0, 1).unsqueeze(0) |
| 44 | for image_ in processed_image |
| 45 | ], dim=0) |
| 46 | return processed_image |
| 47 | |
| 48 | def __call__( |
| 49 | self, |
| 50 | sample, timestep, encoder_hidden_states, conditionings, |
| 51 | tiled=False, tile_size=64, tile_stride=32, **kwargs |
| 52 | ): |
| 53 | res_stack = None |
| 54 | for processor, conditioning, model, scale in zip(self.processors, conditionings, self.models, self.scales): |
| 55 | res_stack_ = model( |
| 56 | sample, timestep, encoder_hidden_states, conditioning, **kwargs, |
| 57 | tiled=tiled, tile_size=tile_size, tile_stride=tile_stride, |
| 58 | processor_id=processor.processor_id |
| 59 | ) |
| 60 | res_stack_ = [res * scale for res in res_stack_] |
| 61 | if res_stack is None: |
| 62 | res_stack = res_stack_ |
| 63 | else: |
| 64 | res_stack = [i + j for i, j in zip(res_stack, res_stack_)] |
| 65 | return res_stack |
| 66 | |
| 67 | |
| 68 | class FluxMultiControlNetManager(MultiControlNetManager): |
no outgoing calls
no test coverage detected