Z-Image-Turbo ControlNet

[hlky]

What does this PR do?

In the original code this is not a typical ControlNet, it is integrated into the transformer and relies on operations performed in the transformer's forward. In this PR we implement it as a typical ControlNet by duplicating the necessary operations from the transformer's forward into the ControlNet's forward and pass transformer to ZImageControlNetModel's forward to access the necessary transformer modules, as a result this is perhaps a little slower than the original implementation, but it keeps things clean and in style. ZImageTransformer2DModel has minimal changes, controlnet_block_samples is introduced, this is a Dict[int, torch.Tensor] returned from ZImageControlNetModel where the int is the ZImageTransformer2DModel layers index, this is another difference from typical ControlNet where every block has the ControlNet output applied. ZImageControlNetPipeline has minimal changes, compared to ZImagePipeline it adds prepare_image function, adds control_image and controlnet_conditioning_scale parameters, prepares and encodes control_image and calls controlnet to obtain controlnet_block_samples which are passed to transformer. control_guidance_start/control_guidance_end is not yet implemented.

Test code

python scripts/convert_z_image_controlnet_to_diffusers.py --original_controlnet_repo_id "alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union" --filename "Z-Image-Turbo-Fun-Controlnet-Union.safetensors" --output_path "z-image-controlnet-hf"

import torch
from diffusers import ZImageControlNetPipeline
from diffusers import ZImageControlNetModel
from diffusers.utils import load_image

controlnet_model = "z-image-controlnet-hf"
controlnet = ZImageControlNetModel.from_pretrained(
    controlnet_model, torch_dtype=torch.bfloat16
)
pipe = ZImageControlNetPipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo", controlnet=controlnet, torch_dtype=torch.bfloat16
)
pipe = pipe.to("cuda")
control_image = load_image("https://huggingface.co/alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union/resolve/main/asset/pose.jpg?download=true")
prompt = "一位年轻女子站在阳光明媚的海岸线上，白裙在轻拂的海风中微微飘动。她拥有一头鲜艳的紫色长发，在风中轻盈舞动，发间系着一个精致的黑色蝴蝶结，与身后柔和的蔚蓝天空形成鲜明对比。她面容清秀，眉目精致，透着一股甜美的青春气息；神情柔和，略带羞涩，目光静静地凝望着远方的地平线，双手自然交叠于身前，仿佛沉浸在思绪之中。在她身后，是辽阔无垠、波光粼粼的大海，阳光洒在海面上，映出温暖的金色光晕。"
image = pipe(
    prompt,
    control_image=control_image,
    controlnet_conditioning_scale=0.75,
    height=1728,
    width=992,
    num_inference_steps=9,
    guidance_scale=0.0,
    generator=torch.Generator("cuda").manual_seed(43),
).images[0]
image.save("zimage.png")

Output

PR	Original

Fixes #12769

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[yiyixuxu]

let's not pass transformer as an input

given that there are some shared layers, we can consider these two alternative design:

• option1: pre-computed shared stuff inside the pipeline, you can add a method to the ZImageTransformer2DModel to be used by controlnet if it makes things easier (but no need to change the transformer code) e.g. inside pipeline

... = self.transformer.prepare_inputs(...)
controlnet_block_samples = self.controlnet(control_image=control_image, ...) 

• Option2: we can try to inject controlnet into transformer inside __init__. similar to https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/animatediff/pipeline_animatediff.py#L140. Basically create a model that combine controlnet + transformer

[hlky]

Which option would you prefer?

[yiyixuxu]

let's try option 2

[hlky]

Done

[elismasilva]

Hi @hlky , I managed to get your PR working locally, but I modified it to work differently. I unified the weights of the base model's transformer with the transformer from the controlnet model. Initially, I did this because I wanted to generate a unified gguf, since when you load the gguf from_single_file you wouldn't be able to load two transformers at the same time. So I unified them and tested with a single gguf and also with a single pretrained model, and it worked. If this is an easier strategy for users, it will only be necessary to upload a different model to Spaces, such as z-image-turbo-control-hf for e.g.

[hlky]

Hi @elismasilva, I see, currently if we used gguf for the transformer and controlnet this would not be applied to the combined model when it's created in the pipeline __init__, I think we can detect the quantization_config and pass that along to the combined model which should fix that, we can also add from_single_file support, a conversion script, etc. if users prefer to have a single weight file, but I did find that injecting the controlnet in __init__ as suggested by @yiyixuxu has increased cpu ram usage from creating the combined transformer+controlnet while the non-combined versions are also loaded, so I wonder if we should go for option 1 instead, we can still add from_single_file support to the controlnet for gguf in that case. Let's wait for some more input from YiYi on how best to proceed.

[elismasilva]

I think the init technique in this case is only valid if the person is going to use the pipeline as a whole. Consider that the ControlNet transformer is not an independent model like those in SDXL, for example; it's an additional weighting of the transformer. So I believe it should be the transformer model's responsibility to append these weights, in case someone tries to load the transformer model in isolation, as is possible, and also considering modular diffusers.

[e1ijah1]

How do we support control modes other than Canny? Other control modes still produce poor results—here’s what I got when I tried the HED example from the demo with prompt “A man holding a bottle” and input image.

[elismasilva]

How do we support control modes other than Canny? Other control modes still produce poor results—here’s what I got when I tried the HED example from the demo with prompt “A man holding a bottle” and input image.

Hi, I ran a test with your image and got the following result:

GT:

HED:

Result 1:
Steps: 9
CFG: 0
Control Scale: 0.7
Prompt: A man holding a bottle

Result 2:
Steps: 9
CFG: 2.5
Control Scale: 0.75
Prompt: raw photo, portrait of a handsome Asian man sitting at a wooden table, holding a green glass bottle, wearing a black sweater, wristwatch, highly detailed skin texture, realistic pores, serious gaze, soft cinematic lighting, rim lighting, balanced exposure, 8k uhd, dslr, sharp focus, wood grain texture.
Negative prompt: underexposed, crushed blacks, too dark, heavy shadows, makeup, smooth skin, plastic, wax, cartoon, illustration, distorted hands, bad anatomy, blur, haze, flat lighting.

To achieve a realistic effect, you will need to apply the Hires.Fix technique to the image after it has been generated:
like this: