Z-Image-Turbo ControlNet #12792

hlky · 2025-12-04T21:31:37Z

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 · 2025-12-05T17:23:18Z

src/diffusers/models/controlnets/controlnet_z_image.py

+
+    def forward(
+        self,
+        transformer: ZImageTransformer2DModel,


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

Which option would you prefer?

let's try option 2

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.

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.

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.

in case someone tries to load the transformer model in isolation, as is possible, and also considering modular diffusers.

are you talking about this?

transformer = ZImageTransformer2DModel.from_pretrained(...) controlnet = ZImageControlNetModel.from_pretrained(...) pipe = ZImageControlNetPipeline.from_pretrained( "Tongyi-MAI/Z-Image-Turbo", controlnet=controlnet, transforfmer=transformer, torch_dtype=torch.bfloat16 )

I believe this PR already handle this use case, no? looking at its __init__

if isinstance(transformer, ZImageTransformer2DModel): transformer = ZImageControlTransformer2DModel.from_controlnet(transformer, controlnet)

the memory situation isn't ideal, but I think we should be able avoid it, I made a few suggestsion in the PR let's test it out to see if it works

in case someone tries to load the transformer model in isolation, as is possible, and also considering modular diffusers.

are you talking about this?

Ah yes, but I had something else in mind; I was thinking about eliminating the "controlnet" object as an intermediary for loading in the transformer's from_controlnet(), I thought that ControlNet could just be another internal module in the same model, but for that to happen we would have to rewrite practically the entire from_pretrained() function to load all the weights in meta device (I believe that's why there's an increase in memory usage during later loading). But anyway, the suggestion in option 3 is also good; you'll remove the things that are currently mixed together, will give a forward() call and a purpose for the controlnet model to exist, but in this case, there would now be two predictions in the pipeline, and the transformer would receive the samples from the controlnet.

e1ijah1 · 2025-12-09T06:42:50Z

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 · 2025-12-09T10:59:50Z

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:

HuggingFaceDocBuilderDev · 2025-12-10T21:48:03Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

yiyixuxu

thanks, I left some comments!

yiyixuxu · 2025-12-10T21:46:37Z

scripts/convert_z_image_controlnet_to_diffusers.py

+    return original_state_dict
+
+
+def convert_z_image_controlnet_checkpoint_to_diffusers(original_state_dict):


there is no conversion going on? maybe we could just load from the original repo using from_single_file so we don't need to host this seperately

yiyixuxu · 2025-12-10T21:47:00Z

src/diffusers/models/controlnets/controlnet_z_image.py

+)
+
+
+class ZImageControlTransformerBlock(ZImageTransformerBlock):


let's not inherit from ZImageTransformerBlock, we can just copy over, modify and make a notes about it

src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py

yiyixuxu · 2025-12-10T22:31:57Z

src/diffusers/models/transformers/transformer_z_image_control.py

+        expected_kwargs, optional_kwargs = cls._get_signature_keys(cls)
+        config = FrozenDict({k: config.get(k) for k in config if k in expected_kwargs or k in optional_kwargs})
+        config["_class_name"] = cls.__name__
+        model = cls.from_config(config)


Suggested change

model = cls.from_config(config)

with init_empty_weights():

model = cls.from_config(config)

yiyixuxu · 2025-12-10T22:35:19Z

src/diffusers/models/transformers/transformer_z_image_control.py

+            return model
+
+        for key, all_x_embedder in transformer.all_x_embedder.items():
+            model.all_x_embedder[key].load_state_dict(all_x_embedder.state_dict())


Can you directly assign these modules instead? This avoids additional memory at loading since we discard the original transformer/controlnet after this anyway.

yiyixuxu · 2025-12-10T22:43:02Z

src/diffusers/models/controlnets/controlnet_z_image.py

+
+    def forward(
+        self,
+        transformer: ZImageTransformer2DModel,


in case someone tries to load the transformer model in isolation, as is possible, and also considering modular diffusers.

are you talking about this?

transformer = ZImageTransformer2DModel.from_pretrained(...) controlnet = ZImageControlNetModel.from_pretrained(...) pipe = ZImageControlNetPipeline.from_pretrained( "Tongyi-MAI/Z-Image-Turbo", controlnet=controlnet, transforfmer=transformer, torch_dtype=torch.bfloat16 )

I believe this PR already handle this use case, no? looking at its __init__

if isinstance(transformer, ZImageTransformer2DModel): transformer = ZImageControlTransformer2DModel.from_controlnet(transformer, controlnet)

the memory situation isn't ideal, but I think we should be able avoid it, I made a few suggestsion in the PR let's test it out to see if it works

src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py

yiyixuxu · 2025-12-11T00:15:08Z

ohm I thought about this a bit more. I think we can should try an Option 3 that's a middle ground between Option 1 and 2: instead of combining everything into one model, the controlnet only load shared layers from the transformer. so your from_transfomer would look lsomething like this

class ZImageControlNetModel:
    @classmethod
    def from_transformer(cls, controlnet, transformer):
        ....
        controlnet.t_embedder = transformer.t_embedder
        controlnet.all_x_embedder = transformer.all_x_embedder
        controlnet.cap_embedder = transformer.cap_embedder
        return controlnet

in pipeline, we still have both controlnet and transformer components and should work similarly to our other contorlnet pipelines

what do you think?

hlky · 2025-12-11T01:25:00Z

@yiyixuxu Option 3 sounds good to me, I've made those changes here a00f104. Let me know if you have any further comments, I will add support for from_single_file next.

hlky · 2025-12-11T03:41:21Z

src/diffusers/loaders/single_file_model.py

+        elif "config_create_fn" in mapping_functions:
+            config_create_fn = mapping_functions["config_create_fn"]
+            diffusers_model_config = config_create_fn()


Normally the config is from the Hub, without hosting the Diffusers format weights this seemed a good option, WDYT? Could be useful in other circumstances, like deriving the config from the weights. Alternative would perhaps be making PR to Hub repo with the config?

hlky · 2025-12-11T03:42:50Z

from_single_file:

import torch
from diffusers import ZImageControlNetModel
from huggingface_hub import hf_hub_download

controlnet = ZImageControlNetModel.from_single_file(
    hf_hub_download(
        "alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union",
        filename="Z-Image-Turbo-Fun-Controlnet-Union.safetensors",
    ),
    torch_dtype=torch.bfloat16,
)

elismasilva · 2025-12-11T04:59:22Z

from_single_file:

import torch
from diffusers import ZImageControlNetModel
from huggingface_hub import hf_hub_download

controlnet = ZImageControlNetModel.from_single_file(
    hf_hub_download(
        "alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union",
        filename="Z-Image-Turbo-Fun-Controlnet-Union.safetensors",
    ),
    torch_dtype=torch.bfloat16,
)

@hlky If you need a ControlNet gguf to test the loading as well, I can try to generate one for you. I'm currently testing this model and trying build an image restoration pipeline, Here is the link to the unified model I mentioned earlier If you want to check it out, go to https://huggingface.co/elismasilva/z-image-control-turbo-unified, but I'll switch to using your implementation later.

yiyixuxu

thanks @hlky
looks good to me, I left some small feedbacks, mainly:

not import class from z-transformer file, it gets messy to import from one model file to another, so let's not do it for now. That being said I think it makes seense to reuse layers in base model for controlnet/ip-adapter and such so we may do that in the future
add missing #Copied from statement

yiyixuxu · 2025-12-11T19:39:31Z

src/diffusers/models/controlnets/controlnet_z_image.py

+from ..controlnets.controlnet import zero_module
+from ..modeling_utils import ModelMixin
+from ..transformers.transformer_z_image import (
+    ADALN_EMBED_DIM,


Can we use #Copied from for these shared layers, instead imports?

at least for now, in the future, we may restructure how we arrange our model code and put the controlnet et all in the same folder as the base model and then maybe we can do it differently

yiyixuxu · 2025-12-11T19:41:47Z

src/diffusers/models/controlnets/controlnet_z_image.py

+        return controlnet
+
+    @staticmethod
+    def create_coordinate_grid(size, start=None, device=None):


can we add a #Coped from?

yiyixuxu · 2025-12-11T19:42:06Z

src/diffusers/models/controlnets/controlnet_z_image.py

+        grids = torch.meshgrid(axes, indexing="ij")
+        return torch.stack(grids, dim=-1)
+
+    def patchify_and_embed(


here too if it's copied from the z-image transformer

yiyixuxu · 2025-12-11T19:49:51Z

cc @DN6 can you take a look for the single file?

hlky added 16 commits December 4, 2025 17:48

init

2354fda

passed transformer

1e2009d

ruff

0c30839

convert passed

52f996e

__init__

4b446b3

pipeline example

a1ff390

ruff

7ab347d

pipeline load_image

8cab0c9

t_scale

8688fa6

x_pad_token

9051272

controlnet_block_samples

0d8c3f1

conditioning_scale

f789325

self.config

5f8ab7b

sample_mode, default controlnet_conditioning_scale

bc72f9c

ruff

13b706a

Merge branch 'main' into z-image-controlnet

728ba02

yiyixuxu reviewed Dec 5, 2025

View reviewed changes

hlky added 5 commits December 5, 2025 20:48

ZImageControlTransformer2DModel

09849a7

ModuleDict

f63a5a8

patchify control_context

f9540cb

transformer weights

3e472ac

-enumerate in ZImageTransformer2DModel

0e7c643

Merge branch 'main' into z-image-controlnet

413c7cb

yiyixuxu reviewed Dec 10, 2025

View reviewed changes

src/diffusers/pipelines/z_image/pipeline_z_image_controlnet.py Show resolved Hide resolved

Option 3

a00f104

from_single_file

a961402

hlky commented Dec 11, 2025

View reviewed changes

Remove convert script

6e1c218

yiyixuxu reviewed Dec 11, 2025

View reviewed changes

yiyixuxu requested a review from DN6 December 11, 2025 19:49

		return original_state_dict


		def convert_z_image_controlnet_checkpoint_to_diffusers(original_state_dict):

		)


		class ZImageControlTransformerBlock(ZImageTransformerBlock):

	model = cls.from_config(config)
	with init_empty_weights():
	model = cls.from_config(config)

Z-Image-Turbo ControlNet #12792

Are you sure you want to change the base?

Z-Image-Turbo ControlNet #12792

Conversation

hlky commented Dec 4, 2025

What does this PR do?

Who can review?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

elismasilva Dec 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

e1ijah1 commented Dec 9, 2025

Uh oh!

elismasilva commented Dec 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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

Uh oh!

HuggingFaceDocBuilderDev commented Dec 10, 2025

Uh oh!

yiyixuxu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yiyixuxu commented Dec 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

hlky commented Dec 11, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hlky commented Dec 11, 2025

Uh oh!

elismasilva commented Dec 11, 2025

Uh oh!

yiyixuxu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yiyixuxu commented Dec 11, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

elismasilva Dec 8, 2025 •

edited

Loading

elismasilva commented Dec 9, 2025 •

edited

Loading

yiyixuxu commented Dec 11, 2025 •

edited

Loading