monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Light refactor of controlnet class

Browse Source
This commit is contained in:
filipstrand 2024-09-22 19:30:21 +02:00
parent 5737b27164
commit d799a924c4
3 changed files with 41 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
src/mflux/controlnet/controlnet_util.py
View File

@ -1,7 +1,11 @@
import logging
import cv2
import numpy as np
import PIL
log = logging.getLogger(__name__)
class ControlnetUtil:
@staticmethod
@ -11,3 +15,10 @@ class ControlnetUtil:
image_to_canny = np.array(image_to_canny[:, :, None])
image_to_canny = np.concatenate([image_to_canny, image_to_canny, image_to_canny], axis=2)
return PIL.Image.fromarray(image_to_canny)
@staticmethod
def scale_image(height: int, width: int, img: PIL.Image) -> PIL.Image:
if height != img.height or width != img.width:
log.warning(f"Control image has different dimensions than the model. Resizing to {width}x{height}")
img = img.resize((width, height), PIL.Image.LANCZOS)
return img

54
src/mflux/controlnet/flux_controlnet.py
View File

@ -1,6 +1,5 @@
import logging
import PIL.Image
import mlx.core as mx
from mlx import nn
from tqdm import tqdm
@ -80,51 +79,51 @@ class Flux1Controlnet:
if weights.quantization_level is not None:
self._set_model_weights(weights)
weights_controlnet, ctrlnet_quantization_level, controlnet_config = WeightHandlerControlnet.load_controlnet_transformer(
weights_controlnet, controlnet_quantization_level, controlnet_config = WeightHandlerControlnet.load_controlnet_transformer(
controlnet_id=CONTROLNET_ID
)
) # fmt: off
self.transformer_controlnet = TransformerControlnet(
model_config=model_config,
num_blocks=controlnet_config["num_layers"],
num_single_blocks=controlnet_config["num_single_layers"],
)
if ctrlnet_quantization_level is None:
if controlnet_quantization_level is None:
self.transformer_controlnet.update(weights_controlnet)
self.bits = None
if quantize is not None or ctrlnet_quantization_level is not None:
self.bits = ctrlnet_quantization_level if ctrlnet_quantization_level is not None else quantize
# fmt: off
nn.quantize(self.transformer_controlnet, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 128, group_size=128, bits=self.bits)
# fmt: on
if quantize is not None or controlnet_quantization_level is not None:
self.bits = controlnet_quantization_level if controlnet_quantization_level is not None else quantize
nn.quantize(self.transformer_controlnet, class_predicate=lambda _, m: isinstance(m, nn.Linear) and len(m.weight[1]) > 128, group_size=128, bits=self.bits) # fmt: off
if ctrlnet_quantization_level is not None:
if controlnet_quantization_level is not None:
self.transformer_controlnet.update(weights_controlnet)
def generate_image(self, seed: int, prompt: str, control_image: PIL.Image.Image, config: ConfigControlnet = ConfigControlnet()) -> GeneratedImage: # fmt: off
def generate_image(
self,
seed: int,
prompt: str,
control_image_path: str,
config: ConfigControlnet = ConfigControlnet()
) -> GeneratedImage: # fmt: off
# Create a new runtime config based on the model type and input parameters
config = RuntimeConfig(config, self.model_config)
time_steps = tqdm(range(config.num_inference_steps))
if config.height != control_image.height or config.width != control_image.width:
log.warning(
f"Control image has different dimensions than the model. Resizing to {config.width}x{config.height}"
)
control_image = control_image.resize((config.width, config.height), PIL.Image.LANCZOS)
# Embedd the controlnet reference image
control_image = ImageUtil.load_image(control_image_path)
control_image = ControlnetUtil.scale_image(config.height, config.width, control_image)
control_image = ControlnetUtil.preprocess_canny(control_image)
controlnet_cond = ImageUtil.to_array(control_image)
controlnet_cond = self.vae.encode(controlnet_cond)
controlnet_cond = (controlnet_cond / self.vae.scaling_factor) + self.vae.shift_factor
controlnet_cond = Flux1Controlnet._pack_latents(controlnet_cond, config.height, config.width)
# 1. Create the initial latents
latents = mx.random.normal(
shape=[1, (config.height // 16) * (config.width // 16), 64],
key=mx.random.key(seed)
) # fmt: off
control_image = ControlnetUtil.preprocess_canny(control_image)
controlnet_cond = ImageUtil.to_array(control_image)
controlnet_cong = self.vae.encode(controlnet_cond)
# the rescaling in the next line is not in the huggingface code, but without it the images from
# the chosen controlnet model are very bad
controlnet_cond = (controlnet_cong / self.vae.scaling_factor) + self.vae.shift_factor
controlnet_cond = Flux1Controlnet._pack_latents(controlnet_cond, config.height, config.width)
# 2. Embedd the prompt
t5_tokens = self.t5_tokenizer.tokenize(prompt)
@ -133,7 +132,8 @@ class Flux1Controlnet:
pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens)
for t in time_steps:
ctrlnet_block_samples, ctrlnet_single_block_samples = self.transformer_controlnet.forward(
# Compute controlnet samples
controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet.forward(
t=t,
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
@ -141,6 +141,7 @@ class Flux1Controlnet:
controlnet_cond=controlnet_cond,
config=config,
)
# 3.t Predict the noise
noise = self.transformer.predict(
t=t,
@ -148,8 +149,8 @@ class Flux1Controlnet:
pooled_prompt_embeds=pooled_prompt_embeds,
hidden_states=latents,
config=config,
controlnet_block_samples=ctrlnet_block_samples,
controlnet_single_block_samples=ctrlnet_single_block_samples,
controlnet_block_samples=controlnet_block_samples,
controlnet_single_block_samples=controlnet_single_block_samples,
)
# 4.t Take one denoise step
@ -171,6 +172,7 @@ class Flux1Controlnet:
lora_paths=self.lora_paths,
lora_scales=self.lora_scales,
config=config,
controlnet_image_path=control_image_path,
)
@staticmethod

4
src/mflux/generate_controlnet.py
View File

@ -1,7 +1,7 @@
import argparse
import time
from mflux import Flux1Controlnet, ConfigControlnet, ModelConfig, ImageUtil
from mflux import Flux1Controlnet, ConfigControlnet, ModelConfig
def main():
@ -45,7 +45,7 @@ def main():
image = flux.generate_image(
seed=int(time.time()) if args.seed is None else args.seed,
prompt=args.prompt,
control_image=ImageUtil.load_image(args.controlnet_image_path),
control_image_path=args.controlnet_image_path,
config=ConfigControlnet(
num_inference_steps=args.steps,
height=args.height,