nodes/aigc-apps-VideoX-Fun
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update Z Image Control 2.1 Lite models (#429)

Browse Source
This commit is contained in:
Bubbliiiing
2026-01-12 21:04:35 +08:00
committed by GitHub
parent a0ba11bc0e
commit 5eb0c6862f
10 changed files with 2679 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
comfyui/z_image/nodes.py
+2
View File
@@ -695,6 +695,7 @@ class LoadZImageControlNetInPipeline:
"required": {
"config": (
[
"z_image/z_image_control_2.1_lite.yaml",
"z_image/z_image_control_2.1.yaml",
"z_image/z_image_control_2.0.yaml",
"z_image/z_image_control_1.0.yaml",
@@ -818,6 +819,7 @@ class LoadZImageControlNetInModel:
"required": {
"config": (
[
"z_image/z_image_control_2.1_lite.yaml",
"z_image/z_image_control_2.1.yaml",
"z_image/z_image_control_2.0.yaml",
"z_image/z_image_control_1.0.yaml",
comfyui/z_image/v1/z_image_chunked_loading_workflow_i2i_inpaint_lite.json
+704
View File
@@ -0,0 +1,704 @@
{
"id": "dcf2fcac-6293-4a86-b30b-f63e420177f2",
"revision": 0,
"last_node_id": 107,
"last_link_id": 113,
"nodes": [
{
"id": 78,
"type": "Note",
"pos": [
18,
-46
],
"size": [
210,
88
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [],
"properties": {
"text": ""
},
"widgets_values": [
"You can write prompt here\n（你可以在此填写提示词）"
],
"color": "#432",
"bgcolor": "#653"
},
{
"id": 91,
"type": "LoadZImageTextEncoderModel",
"pos": [
283.53765869140625,
-280.6837463378906
],
"size": [
407.4130859375,
102
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "text_encoder",
"type": "TextEncoderModel",
"links": [
80
]
},
{
"name": "tokenizer",
"type": "Tokenizer",
"links": [
81
]
}
],
"properties": {
"Node name for S&R": "LoadZImageTextEncoderModel"
},
"widgets_values": [
"qwen_3_4b.safetensors",
"bf16"
]
},
{
"id": 75,
"type": "FunTextBox",
"pos": [
250,
-50
],
"size": [
383.54010009765625,
156.71620178222656
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "prompt",
"type": "STRING_PROMPT",
"slot_index": 0,
"links": [
88
]
}
],
"title": "Positive Prompt（正向提示词）",
"properties": {
"Node name for S&R": "FunTextBox"
},
"widgets_values": [
"A photo of Sakura, a 17-year-old high school student from Japan, captured in a candid, high-fidelity cinematic moment on a rainy evening. She is squatting low on the rain-slicked asphalt of an urban sidewalk, holding a transparent vinyl umbrella with a white handle resting over her shoulder in one hand, her other hand resting on her knee. The clear plastic canopy is streaked with rivulets of water and beaded with droplets that catch the ambient city light. A profound, silent interaction defines the scene: Sakura is looking directly downward, her expression gentle and focused, locking eyes with a small black cat sitting on the wet ground in front of her.\n\nSakura has long, lustrous black hair styled in a precise hime cut with blunt bangs across her forehead and sidelocks framing her cheeks, damp strands clinging subtly to her jacket, with a single red ribbon tied on the left side. Her visible pores on her nose, and a soft sheen of moisture on her cheeks. She wears a dark navy sailor-style school uniform (seifuku) featuring a white collar with red linear detailing and a bright red necktie loosely knotted at the chest; a simple black choker encircles her neck. The uniform jacket has oversized sleeves. Her lower body features a short, dark pleated miniskirt that fans slightly over clean white ankle socks that provide a stark contrast to the wet asphalt, ending in dark leather loafers that gleam with moisture.\n\nThe black cat sits upright in a shallow puddle, its short fur slicked by the rain, tilting its head back to stare intently up into Sakura's face, establishing a clear line of sight. The background is anchored by a large, illuminated red vending machine standing against the darkness, its cool bluish-white interior light spilling onto Sakura's profile and the umbrella. The ground reflects the red chassis and the neon streetlights in distorted patches on the wet pavement. Additional cool rain streaks fall through the frame, some caught in sharp focus and others blurred into vertical lines against the background lights. The scene is rendered with a wide-aperture lens creating a shallow depth of field, keeping the girl and cat in sharp focus while softening the background into gentle bokeh, with the texture of fine-grain 35mm film stock.\n"
]
},
{
"id": 73,
"type": "FunTextBox",
"pos": [
250,
160
],
"size": [
383.7149963378906,
183.83506774902344
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "prompt",
"type": "STRING_PROMPT",
"slot_index": 0,
"links": [
89
]
}
],
"title": "Negtive Prompt（反向提示词）",
"properties": {
"Node name for S&R": "FunTextBox"
},
"widgets_values": [
""
]
},
{
"id": 97,
"type": "Note",
"pos": [
-354.4680507215508,
-433.6570714778354
],
"size": [
598.1623727144193,
233.48501180428053
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [],
"outputs": [],
"properties": {
"text": ""
},
"widgets_values": [
"GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].\nmodel_full_load means that the entire model will be moved to the GPU.\n\nmodel_full_load_and_qfloat8 means that the entire model will be moved to the GPU,\nand the transformer model has been quantized to float8, which can save more GPU memory. \n\nmodel_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory.\n\nmodel_cpu_offload_and_qfloat8 indicates that the entire model will be moved to the CPU after use, \nand the transformer model has been quantized to float8, which can save more GPU memory. \n\nsequential_cpu_offload means that each layer of the model will be moved to the CPU after use, \nresulting in slower speeds but saving a large amount of GPU memory."
],
"color": "#432",
"bgcolor": "#653"
},
{
"id": 96,
"type": "CombineZImagePipeline",
"pos": [
790.3572998046875,
-328.7134094238281
],
"size": [
342.5804748535156,
162
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "transformer",
"type": "TransformerModel",
"link": 96
},
{
"name": "vae",
"type": "VAEModel",
"link": 79
},
{
"name": "text_encoder",
"type": "TextEncoderModel",
"link": 80
},
{
"name": "tokenizer",
"type": "Tokenizer",
"link": 81
},
{
"name": "processor",
"shape": 7,
"type": "Processor",
"link": null
},
{
"name": "model_name",
"type": "STRING",
"widget": {
"name": "model_name"
},
"link": 82
}
],
"outputs": [
{
"name": "funmodels",
"type": "FunModels",
"links": [
87
]
}
],
"properties": {
"Node name for S&R": "CombineZImagePipeline"
},
"widgets_values": [
"",
"model_cpu_offload"
]
},
{
"id": 92,
"type": "LoadZImageTransformerModel",
"pos": [
275.9798278808594,
-465.2391052246094
],
"size": [
416.3677673339844,
106.13789367675781
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "transformer",
"type": "TransformerModel",
"links": [
95
]
},
{
"name": "model_name",
"type": "STRING",
"links": [
82
]
}
],
"properties": {
"Node name for S&R": "LoadZImageTransformerModel"
},
"widgets_values": [
"z_image_turbo_bf16.safetensors",
"bf16"
]
},
{
"id": 99,
"type": "ZImageControlSampler",
"pos": [
727.2482831521481,
-52.41010988674983
],
"size": [
270,
350
],
"flags": {},
"order": 12,
"mode": 0,
"inputs": [
{
"name": "funmodels",
"type": "FunModels",
"link": 87
},
{
"name": "prompt",
"type": "STRING_PROMPT",
"link": 88
},
{
"name": "negative_prompt",
"type": "STRING_PROMPT",
"link": 89
},
{
"name": "control_image",
"shape": 7,
"type": "IMAGE",
"link": null
},
{
"name": "inpaint_image",
"shape": 7,
"type": "IMAGE",
"link": 110
},
{
"name": "mask_image",
"shape": 7,
"type": "IMAGE",
"link": 112
}
],
"outputs": [
{
"name": "images",
"type": "IMAGE",
"links": [
90
]
}
],
"properties": {
"Node name for S&R": "ZImageControlSampler"
},
"widgets_values": [
1184,
1568,
43,
"fixed",
8,
0,
"Flow",
3,
0.8
]
},
{
"id": 93,
"type": "LoadZImageVAEModel",
"pos": [
1168.1599508804559,
-314.8650476692169
],
"size": [
377.8583984375,
84.69844055175781
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "vae",
"type": "VAEModel",
"links": [
79
]
}
],
"properties": {
"Node name for S&R": "LoadZImageVAEModel"
},
"widgets_values": [
"ae.safetensors",
"bf16"
]
},
{
"id": 88,
"type": "PreviewImage",
"pos": [
1049.1402001998824,
-52.699751832945104
],
"size": [
366.56134033203125,
415.4429626464844
],
"flags": {},
"order": 13,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 90
}
],
"outputs": [],
"properties": {
"Node name for S&R": "PreviewImage"
},
"widgets_values": []
},
{
"id": 104,
"type": "PreviewImage",
"pos": [
911.4377073728218,
418.9988584275326
],
"size": [
366.56134033203125,
415.4429626464844
],
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 113
}
],
"outputs": [],
"properties": {
"Node name for S&R": "PreviewImage"
},
"widgets_values": []
},
{
"id": 107,
"type": "MaskToImage",
"pos": [
658.3044275669289,
511.8717365608471
],
"size": [
140,
26
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "mask",
"type": "MASK",
"link": 111
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
112,
113
]
}
],
"properties": {
"Node name for S&R": "MaskToImage"
},
"widgets_values": []
},
{
"id": 100,
"type": "LoadImage",
"pos": [
323.5443519405259,
489.8702544908114
],
"size": [
270,
314.00000000000006
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
110
]
},
{
"name": "MASK",
"type": "MASK",
"links": [
111
]
}
],
"properties": {
"Node name for S&R": "LoadImage",
"image": "clipspace/clipspace-painted-masked-1766731857414.png [input]"
},
"widgets_values": [
"clipspace/clipspace-painted-masked-1766731857414.png [input]",
"image"
]
},
{
"id": 102,
"type": "LoadZImageControlNetInModel",
"pos": [
779.793189390101,
-457.3825558553134
],
"size": [
589.8698159570357,
82
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "transformer",
"type": "TransformerModel",
"link": 95
}
],
"outputs": [
{
"name": "transformer",
"type": "TransformerModel",
"links": [
96
]
}
],
"properties": {
"Node name for S&R": "LoadZImageControlNetInModel"
},
"widgets_values": [
"z_image/z_image_control_2.1_lite.yaml",
"Z-Image-Turbo-Fun-Controlnet-Union-2.1-lite-2601-8steps.safetensors"
]
}
],
"links": [
[
79,
93,
0,
96,
1,
"VAEModel"
],
[
80,
91,
0,
96,
2,
"TextEncoderModel"
],
[
81,
91,
1,
96,
3,
"Tokenizer"
],
[
82,
92,
1,
96,
5,
"STRING"
],
[
87,
96,
0,
99,
0,
"FunModels"
],
[
88,
75,
0,
99,
1,
"STRING_PROMPT"
],
[
89,
73,
0,
99,
2,
"STRING_PROMPT"
],
[
90,
99,
0,
88,
0,
"IMAGE"
],
[
95,
92,
0,
102,
0,
"TransformerModel"
],
[
96,
102,
0,
96,
0,
"TransformerModel"
],
[
110,
100,
0,
99,
4,
"IMAGE"
],
[
111,
100,
1,
107,
0,
"MASK"
],
[
112,
107,
0,
99,
5,
"IMAGE"
],
[
113,
107,
0,
104,
0,
"IMAGE"
]
],
"groups": [
{
"id": 1,
"title": "Load Model",
"bounding": [
227.96267700195312,
-546.4359741210938,
1350.4793699732413,
404.87677206390265
],
"color": "#b06634",
"font_size": 24,
"flags": {}
},
{
"id": 2,
"title": "Prompts",
"bounding": [
218,
-127,
450,
483
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
}
],
"config": {},
"extra": {
"ds": {
"scale": 0.7772383863288174,
"offset": [
248.90036994067935,
733.942567031054
]
},
"frontendVersion": "1.34.9",
"workflowRendererVersion": "LG",
"workspace_info": {
"id": "776b62b4-bd17-4ed3-9923-b7aad000b1ea"
},
"node_versions": {
"CogVideoX-Fun": "07dd34b942f866d5f95e8b812b6082d359079260",
"comfy-core": "0.6.0"
}
},
"version": 0.4
}
comfyui/z_image/v1/z_image_chunked_loading_workflow_t2i_control_lite.json
+614
View File
@@ -0,0 +1,614 @@
{
"id": "dcf2fcac-6293-4a86-b30b-f63e420177f2",
"revision": 0,
"last_node_id": 102,
"last_link_id": 96,
"nodes": [
{
"id": 78,
"type": "Note",
"pos": [
18,
-46
],
"size": [
210,
88
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [],
"properties": {
"text": ""
},
"widgets_values": [
"You can write prompt here\n（你可以在此填写提示词）"
],
"color": "#432",
"bgcolor": "#653"
},
{
"id": 91,
"type": "LoadZImageTextEncoderModel",
"pos": [
283.53765869140625,
-280.6837463378906
],
"size": [
407.4130859375,
102
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "text_encoder",
"type": "TextEncoderModel",
"links": [
80
]
},
{
"name": "tokenizer",
"type": "Tokenizer",
"links": [
81
]
}
],
"properties": {
"Node name for S&R": "LoadZImageTextEncoderModel"
},
"widgets_values": [
"qwen_3_4b.safetensors",
"bf16"
]
},
{
"id": 75,
"type": "FunTextBox",
"pos": [
250,
-50
],
"size": [
383.54010009765625,
156.71620178222656
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "prompt",
"type": "STRING_PROMPT",
"slot_index": 0,
"links": [
88
]
}
],
"title": "Positive Prompt（正向提示词）",
"properties": {
"Node name for S&R": "FunTextBox"
},
"widgets_values": [
"A photo of Sakura, a 17-year-old high school student from Japan, captured in a candid, high-fidelity cinematic moment on a rainy evening. She is squatting low on the rain-slicked asphalt of an urban sidewalk, holding a transparent vinyl umbrella with a white handle resting over her shoulder in one hand, her other hand resting on her knee. The clear plastic canopy is streaked with rivulets of water and beaded with droplets that catch the ambient city light. A profound, silent interaction defines the scene: Sakura is looking directly downward, her expression gentle and focused, locking eyes with a small black cat sitting on the wet ground in front of her.\n\nSakura has long, lustrous black hair styled in a precise hime cut with blunt bangs across her forehead and sidelocks framing her cheeks, damp strands clinging subtly to her jacket, with a single red ribbon tied on the left side. Her visible pores on her nose, and a soft sheen of moisture on her cheeks. She wears a dark navy sailor-style school uniform (seifuku) featuring a white collar with red linear detailing and a bright red necktie loosely knotted at the chest; a simple black choker encircles her neck. The uniform jacket has oversized sleeves. Her lower body features a short, dark pleated miniskirt that fans slightly over clean white ankle socks that provide a stark contrast to the wet asphalt, ending in dark leather loafers that gleam with moisture.\n\nThe black cat sits upright in a shallow puddle, its short fur slicked by the rain, tilting its head back to stare intently up into Sakura's face, establishing a clear line of sight. The background is anchored by a large, illuminated red vending machine standing against the darkness, its cool bluish-white interior light spilling onto Sakura's profile and the umbrella. The ground reflects the red chassis and the neon streetlights in distorted patches on the wet pavement. Additional cool rain streaks fall through the frame, some caught in sharp focus and others blurred into vertical lines against the background lights. The scene is rendered with a wide-aperture lens creating a shallow depth of field, keeping the girl and cat in sharp focus while softening the background into gentle bokeh, with the texture of fine-grain 35mm film stock.\n"
]
},
{
"id": 73,
"type": "FunTextBox",
"pos": [
250,
160
],
"size": [
383.7149963378906,
183.83506774902344
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "prompt",
"type": "STRING_PROMPT",
"slot_index": 0,
"links": [
89
]
}
],
"title": "Negtive Prompt（反向提示词）",
"properties": {
"Node name for S&R": "FunTextBox"
},
"widgets_values": [
""
]
},
{
"id": 97,
"type": "Note",
"pos": [
-354.4680507215508,
-433.6570714778354
],
"size": [
598.1623727144193,
233.48501180428053
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [],
"outputs": [],
"properties": {
"text": ""
},
"widgets_values": [
"GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].\nmodel_full_load means that the entire model will be moved to the GPU.\n\nmodel_full_load_and_qfloat8 means that the entire model will be moved to the GPU,\nand the transformer model has been quantized to float8, which can save more GPU memory. \n\nmodel_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory.\n\nmodel_cpu_offload_and_qfloat8 indicates that the entire model will be moved to the CPU after use, \nand the transformer model has been quantized to float8, which can save more GPU memory. \n\nsequential_cpu_offload means that each layer of the model will be moved to the CPU after use, \nresulting in slower speeds but saving a large amount of GPU memory."
],
"color": "#432",
"bgcolor": "#653"
},
{
"id": 93,
"type": "LoadZImageVAEModel",
"pos": [
1168.1599508804559,
-314.8650476692169
],
"size": [
377.8583984375,
84.69844055175781
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "vae",
"type": "VAEModel",
"links": [
79
]
}
],
"properties": {
"Node name for S&R": "LoadZImageVAEModel"
},
"widgets_values": [
"ae.safetensors",
"bf16"
]
},
{
"id": 88,
"type": "PreviewImage",
"pos": [
1049.1402001998824,
-52.699751832945104
],
"size": [
366.56134033203125,
415.4429626464844
],
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 90
}
],
"outputs": [],
"properties": {
"Node name for S&R": "PreviewImage"
},
"widgets_values": []
},
{
"id": 100,
"type": "LoadImage",
"pos": [
396.48551767952716,
419.158511044569
],
"size": [
270,
314.00000000000006
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
86
]
},
{
"name": "MASK",
"type": "MASK",
"links": null
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"a7kXeQ5l9Dhspes7q3x3G (1).png",
"image"
]
},
{
"id": 92,
"type": "LoadZImageTransformerModel",
"pos": [
275.9798278808594,
-465.2391052246094
],
"size": [
416.3677673339844,
106.13789367675781
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "transformer",
"type": "TransformerModel",
"links": [
95
]
},
{
"name": "model_name",
"type": "STRING",
"links": [
82
]
}
],
"properties": {
"Node name for S&R": "LoadZImageTransformerModel"
},
"widgets_values": [
"z_image_turbo_bf16.safetensors",
"bf16"
]
},
{
"id": 99,
"type": "ZImageControlSampler",
"pos": [
727.2482831521481,
-52.41010988674983
],
"size": [
270,
350
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "funmodels",
"type": "FunModels",
"link": 87
},
{
"name": "prompt",
"type": "STRING_PROMPT",
"link": 88
},
{
"name": "negative_prompt",
"type": "STRING_PROMPT",
"link": 89
},
{
"name": "control_image",
"shape": 7,
"type": "IMAGE",
"link": 86
},
{
"name": "inpaint_image",
"shape": 7,
"type": "IMAGE",
"link": null
},
{
"name": "mask_image",
"shape": 7,
"type": "IMAGE",
"link": null
}
],
"outputs": [
{
"name": "images",
"type": "IMAGE",
"links": [
90
]
}
],
"properties": {
"Node name for S&R": "ZImageControlSampler"
},
"widgets_values": [
1184,
1568,
43,
"fixed",
8,
0,
"Flow",
3,
0.8
]
},
{
"id": 96,
"type": "CombineZImagePipeline",
"pos": [
790.3572998046875,
-328.7134094238281
],
"size": [
342.5804748535156,
162
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "transformer",
"type": "TransformerModel",
"link": 96
},
{
"name": "vae",
"type": "VAEModel",
"link": 79
},
{
"name": "text_encoder",
"type": "TextEncoderModel",
"link": 80
},
{
"name": "tokenizer",
"type": "Tokenizer",
"link": 81
},
{
"name": "processor",
"shape": 7,
"type": "Processor",
"link": null
},
{
"name": "model_name",
"type": "STRING",
"widget": {
"name": "model_name"
},
"link": 82
}
],
"outputs": [
{
"name": "funmodels",
"type": "FunModels",
"links": [
87
]
}
],
"properties": {
"Node name for S&R": "CombineZImagePipeline"
},
"widgets_values": [
"",
"model_cpu_offload"
]
},
{
"id": 102,
"type": "LoadZImageControlNetInModel",
"pos": [
779.793189390101,
-457.3825558553134
],
"size": [
589.8698159570357,
82
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "transformer",
"type": "TransformerModel",
"link": 95
}
],
"outputs": [
{
"name": "transformer",
"type": "TransformerModel",
"links": [
96
]
}
],
"properties": {
"Node name for S&R": "LoadZImageControlNetInModel"
},
"widgets_values": [
"z_image/z_image_control_2.1_lite.yaml",
"Z-Image-Turbo-Fun-Controlnet-Union-2.1-lite-2601-8steps.safetensors"
]
}
],
"links": [
[
79,
93,
0,
96,
1,
"VAEModel"
],
[
80,
91,
0,
96,
2,
"TextEncoderModel"
],
[
81,
91,
1,
96,
3,
"Tokenizer"
],
[
82,
92,
1,
96,
5,
"STRING"
],
[
86,
100,
0,
99,
3,
"IMAGE"
],
[
87,
96,
0,
99,
0,
"FunModels"
],
[
88,
75,
0,
99,
1,
"STRING_PROMPT"
],
[
89,
73,
0,
99,
2,
"STRING_PROMPT"
],
[
90,
99,
0,
88,
0,
"IMAGE"
],
[
95,
92,
0,
102,
0,
"TransformerModel"
],
[
96,
102,
0,
96,
0,
"TransformerModel"
]
],
"groups": [
{
"id": 1,
"title": "Load Model",
"bounding": [
227.96267700195312,
-546.4359741210938,
1350.4793699732413,
404.87677206390265
],
"color": "#b06634",
"font_size": 24,
"flags": {}
},
{
"id": 2,
"title": "Prompts",
"bounding": [
218,
-127,
450,
483
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
}
],
"config": {},
"extra": {
"ds": {
"scale": 0.7208430239838531,
"offset": [
460.4224504078358,
644.7360602102879
]
},
"frontendVersion": "1.34.9",
"workflowRendererVersion": "LG",
"workspace_info": {
"id": "776b62b4-bd17-4ed3-9923-b7aad000b1ea"
},
"node_versions": {
"CogVideoX-Fun": "07dd34b942f866d5f95e8b812b6082d359079260",
"comfy-core": "0.6.0"
}
},
"version": 0.4
}
comfyui/z_image/v1/z_image_chunked_loading_workflow_tile_lite.json
+614
View File
@@ -0,0 +1,614 @@
{
"id": "dcf2fcac-6293-4a86-b30b-f63e420177f2",
"revision": 0,
"last_node_id": 102,
"last_link_id": 96,
"nodes": [
{
"id": 78,
"type": "Note",
"pos": [
18,
-46
],
"size": [
210,
88
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [],
"properties": {
"text": ""
},
"widgets_values": [
"You can write prompt here\n（你可以在此填写提示词）"
],
"color": "#432",
"bgcolor": "#653"
},
{
"id": 91,
"type": "LoadZImageTextEncoderModel",
"pos": [
283.53765869140625,
-280.6837463378906
],
"size": [
407.4130859375,
102
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "text_encoder",
"type": "TextEncoderModel",
"links": [
80
]
},
{
"name": "tokenizer",
"type": "Tokenizer",
"links": [
81
]
}
],
"properties": {
"Node name for S&R": "LoadZImageTextEncoderModel"
},
"widgets_values": [
"qwen_3_4b.safetensors",
"bf16"
]
},
{
"id": 75,
"type": "FunTextBox",
"pos": [
250,
-50
],
"size": [
383.54010009765625,
156.71620178222656
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "prompt",
"type": "STRING_PROMPT",
"slot_index": 0,
"links": [
88
]
}
],
"title": "Positive Prompt（正向提示词）",
"properties": {
"Node name for S&R": "FunTextBox"
},
"widgets_values": [
"A photo of Sakura, a 17-year-old high school student from Japan, captured in a candid, high-fidelity cinematic moment on a rainy evening. She is squatting low on the rain-slicked asphalt of an urban sidewalk, holding a transparent vinyl umbrella with a white handle resting over her shoulder in one hand, her other hand resting on her knee. The clear plastic canopy is streaked with rivulets of water and beaded with droplets that catch the ambient city light. A profound, silent interaction defines the scene: Sakura is looking directly downward, her expression gentle and focused, locking eyes with a small black cat sitting on the wet ground in front of her.\n\nSakura has long, lustrous black hair styled in a precise hime cut with blunt bangs across her forehead and sidelocks framing her cheeks, damp strands clinging subtly to her jacket, with a single red ribbon tied on the left side. Her visible pores on her nose, and a soft sheen of moisture on her cheeks. She wears a dark navy sailor-style school uniform (seifuku) featuring a white collar with red linear detailing and a bright red necktie loosely knotted at the chest; a simple black choker encircles her neck. The uniform jacket has oversized sleeves. Her lower body features a short, dark pleated miniskirt that fans slightly over clean white ankle socks that provide a stark contrast to the wet asphalt, ending in dark leather loafers that gleam with moisture.\n\nThe black cat sits upright in a shallow puddle, its short fur slicked by the rain, tilting its head back to stare intently up into Sakura's face, establishing a clear line of sight. The background is anchored by a large, illuminated red vending machine standing against the darkness, its cool bluish-white interior light spilling onto Sakura's profile and the umbrella. The ground reflects the red chassis and the neon streetlights in distorted patches on the wet pavement. Additional cool rain streaks fall through the frame, some caught in sharp focus and others blurred into vertical lines against the background lights. The scene is rendered with a wide-aperture lens creating a shallow depth of field, keeping the girl and cat in sharp focus while softening the background into gentle bokeh, with the texture of fine-grain 35mm film stock.\n"
]
},
{
"id": 73,
"type": "FunTextBox",
"pos": [
250,
160
],
"size": [
383.7149963378906,
183.83506774902344
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "prompt",
"type": "STRING_PROMPT",
"slot_index": 0,
"links": [
89
]
}
],
"title": "Negtive Prompt（反向提示词）",
"properties": {
"Node name for S&R": "FunTextBox"
},
"widgets_values": [
""
]
},
{
"id": 97,
"type": "Note",
"pos": [
-354.4680507215508,
-433.6570714778354
],
"size": [
598.1623727144193,
233.48501180428053
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [],
"outputs": [],
"properties": {
"text": ""
},
"widgets_values": [
"GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].\nmodel_full_load means that the entire model will be moved to the GPU.\n\nmodel_full_load_and_qfloat8 means that the entire model will be moved to the GPU,\nand the transformer model has been quantized to float8, which can save more GPU memory. \n\nmodel_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory.\n\nmodel_cpu_offload_and_qfloat8 indicates that the entire model will be moved to the CPU after use, \nand the transformer model has been quantized to float8, which can save more GPU memory. \n\nsequential_cpu_offload means that each layer of the model will be moved to the CPU after use, \nresulting in slower speeds but saving a large amount of GPU memory."
],
"color": "#432",
"bgcolor": "#653"
},
{
"id": 93,
"type": "LoadZImageVAEModel",
"pos": [
1168.1599508804559,
-314.8650476692169
],
"size": [
377.8583984375,
84.69844055175781
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "vae",
"type": "VAEModel",
"links": [
79
]
}
],
"properties": {
"Node name for S&R": "LoadZImageVAEModel"
},
"widgets_values": [
"ae.safetensors",
"bf16"
]
},
{
"id": 88,
"type": "PreviewImage",
"pos": [
1049.1402001998824,
-52.699751832945104
],
"size": [
366.56134033203125,
415.4429626464844
],
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 90
}
],
"outputs": [],
"properties": {
"Node name for S&R": "PreviewImage"
},
"widgets_values": []
},
{
"id": 100,
"type": "LoadImage",
"pos": [
396.48551767952716,
419.158511044569
],
"size": [
270,
314.00000000000006
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
86
]
},
{
"name": "MASK",
"type": "MASK",
"links": null
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"z-image-turbo_00004_.png",
"image"
]
},
{
"id": 92,
"type": "LoadZImageTransformerModel",
"pos": [
275.9798278808594,
-465.2391052246094
],
"size": [
416.3677673339844,
106.13789367675781
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "transformer",
"type": "TransformerModel",
"links": [
95
]
},
{
"name": "model_name",
"type": "STRING",
"links": [
82
]
}
],
"properties": {
"Node name for S&R": "LoadZImageTransformerModel"
},
"widgets_values": [
"z_image_turbo_bf16.safetensors",
"bf16"
]
},
{
"id": 99,
"type": "ZImageControlSampler",
"pos": [
727.2482831521481,
-52.41010988674983
],
"size": [
270,
350
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "funmodels",
"type": "FunModels",
"link": 87
},
{
"name": "prompt",
"type": "STRING_PROMPT",
"link": 88
},
{
"name": "negative_prompt",
"type": "STRING_PROMPT",
"link": 89
},
{
"name": "control_image",
"shape": 7,
"type": "IMAGE",
"link": 86
},
{
"name": "inpaint_image",
"shape": 7,
"type": "IMAGE",
"link": null
},
{
"name": "mask_image",
"shape": 7,
"type": "IMAGE",
"link": null
}
],
"outputs": [
{
"name": "images",
"type": "IMAGE",
"links": [
90
]
}
],
"properties": {
"Node name for S&R": "ZImageControlSampler"
},
"widgets_values": [
1824,
2416,
43,
"fixed",
8,
0,
"Flow",
3,
0.8
]
},
{
"id": 102,
"type": "LoadZImageControlNetInModel",
"pos": [
779.793189390101,
-457.3825558553134
],
"size": [
589.8698159570357,
82
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "transformer",
"type": "TransformerModel",
"link": 95
}
],
"outputs": [
{
"name": "transformer",
"type": "TransformerModel",
"links": [
96
]
}
],
"properties": {
"Node name for S&R": "LoadZImageControlNetInModel"
},
"widgets_values": [
"z_image/z_image_control_2.1_lite.yaml",
"Z-Image-Turbo-Fun-Controlnet-Tile-2.1-lite-2601-8steps.safetensors"
]
},
{
"id": 96,
"type": "CombineZImagePipeline",
"pos": [
790.3572998046875,
-328.7134094238281
],
"size": [
342.5804748535156,
162
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "transformer",
"type": "TransformerModel",
"link": 96
},
{
"name": "vae",
"type": "VAEModel",
"link": 79
},
{
"name": "text_encoder",
"type": "TextEncoderModel",
"link": 80
},
{
"name": "tokenizer",
"type": "Tokenizer",
"link": 81
},
{
"name": "processor",
"shape": 7,
"type": "Processor",
"link": null
},
{
"name": "model_name",
"type": "STRING",
"widget": {
"name": "model_name"
},
"link": 82
}
],
"outputs": [
{
"name": "funmodels",
"type": "FunModels",
"links": [
87
]
}
],
"properties": {
"Node name for S&R": "CombineZImagePipeline"
},
"widgets_values": [
"",
"sequential_cpu_offload"
]
}
],
"links": [
[
79,
93,
0,
96,
1,
"VAEModel"
],
[
80,
91,
0,
96,
2,
"TextEncoderModel"
],
[
81,
91,
1,
96,
3,
"Tokenizer"
],
[
82,
92,
1,
96,
5,
"STRING"
],
[
86,
100,
0,
99,
3,
"IMAGE"
],
[
87,
96,
0,
99,
0,
"FunModels"
],
[
88,
75,
0,
99,
1,
"STRING_PROMPT"
],
[
89,
73,
0,
99,
2,
"STRING_PROMPT"
],
[
90,
99,
0,
88,
0,
"IMAGE"
],
[
95,
92,
0,
102,
0,
"TransformerModel"
],
[
96,
102,
0,
96,
0,
"TransformerModel"
]
],
"groups": [
{
"id": 1,
"title": "Load Model",
"bounding": [
227.96267700195312,
-546.4359741210938,
1350.4793699732413,
404.87677206390265
],
"color": "#b06634",
"font_size": 24,
"flags": {}
},
{
"id": 2,
"title": "Prompts",
"bounding": [
218,
-127,
450,
483
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
}
],
"config": {},
"extra": {
"ds": {
"scale": 0.6477940671634007,
"offset": [
475.03594149909674,
811.7170336004714
]
},
"frontendVersion": "1.34.9",
"workflowRendererVersion": "LG",
"workspace_info": {
"id": "776b62b4-bd17-4ed3-9923-b7aad000b1ea"
},
"node_versions": {
"CogVideoX-Fun": "07dd34b942f866d5f95e8b812b6082d359079260",
"comfy-core": "0.6.0"
}
},
"version": 0.4
}
config/z_image/z_image_control_2.1_lite.yaml
+8
View File
@@ -0,0 +1,8 @@
format: diffusers
pipeline: z_image
transformer_additional_kwargs:
control_layers_places: [0, 10, 20]
control_refiner_layers_places: [0, 1]
add_control_noise_refiner: true
add_control_noise_refiner_correctly: true
control_in_dim: 33
examples/z_image_fun/predict_i2i_inpaint_2.1_lite.py
+241
View File
@@ -0,0 +1,241 @@
import os
import sys
import numpy as np
import torch
from diffusers import FlowMatchEulerDiscreteScheduler
from omegaconf import OmegaConf
from PIL import Image
current_file_path = os.path.abspath(__file__)
project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
for project_root in project_roots:
sys.path.insert(0, project_root) if project_root not in sys.path else None
from videox_fun.dist import set_multi_gpus_devices, shard_model
from videox_fun.models import (AutoencoderKL, AutoTokenizer,
Qwen3ForCausalLM, ZImageControlTransformer2DModel)
from videox_fun.models.cache_utils import get_teacache_coefficients
from videox_fun.pipeline import ZImageControlPipeline
from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
convert_weight_dtype_wrapper)
from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent, get_image_latent, get_image,
get_video_to_video_latent,
save_videos_grid)
# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
# model_full_load means that the entire model will be moved to the GPU.
#
# model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
# and the transformer model has been quantized to float8, which can save more GPU memory.
#
# model_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory.
#
# model_cpu_offload_and_qfloat8 indicates that the entire model will be moved to the CPU after use,
# and the transformer model has been quantized to float8, which can save more GPU memory.
#
# sequential_cpu_offload means that each layer of the model will be moved to the CPU after use,
# resulting in slower speeds but saving a large amount of GPU memory.
GPU_memory_mode = "model_cpu_offload"
# Multi GPUs config
# Please ensure that the product of ulysses_degree and ring_degree equals the number of GPUs used.
# For example, if you are using 8 GPUs, you can set ulysses_degree = 2 and ring_degree = 4.
# If you are using 1 GPU, you can set ulysses_degree = 1 and ring_degree = 1.
ulysses_degree = 1
ring_degree = 1
# Use FSDP to save more GPU memory in multi gpus.
fsdp_dit = False
fsdp_text_encoder = False
# Compile will give a speedup in fixed resolution and need a little GPU memory.
# The compile_dit is not compatible with the fsdp_dit and sequential_cpu_offload.
compile_dit = False
# Config and model path
config_path = "config/z_image/z_image_control_2.1_lite.yaml"
# model path
model_name = "models/Diffusion_Transformer/Z-Image-Turbo"
# Choose the sampler in "Flow", "Flow_Unipc", "Flow_DPM++"
sampler_name = "Flow"
# Load pretrained model if need
transformer_path = "models/Personalized_Model/Z-Image-Turbo-Fun-Controlnet-Union-2.1-lite-2601-8steps.safetensors"
vae_path = None
lora_path = None
# Other params
sample_size = [1728, 992]
# Use torch.float16 if GPU does not support torch.bfloat16
# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
weight_dtype = torch.bfloat16
control_image = "asset/pose.jpg"
inpaint_image = "asset/8.png"
mask_image = "asset/mask.png"
control_context_scale = 0.85
# Please use as detailed a prompt as possible to describe the object that needs to be generated.
prompt = "画面中央是一位年轻女孩，她拥有一头令人印象深刻的亮紫色长发，发丝在海风中轻盈飘扬，营造出动感而唯美的效果。她的长发两侧各扎着黑色蝴蝶结发饰，增添了几分可爱与俏皮感。女孩身穿一袭纯白色无袖连衣裙，裙摆轻盈飘逸，与她清新的气质完美契合。她的妆容精致自然，淡粉色的唇妆和温柔的眼神流露出恬静优雅的气质。她单手叉腰，姿态自信从容，目光直视镜头，展现出既甜美又不失个性的魅力。背景是一片开阔的海景，湛蓝的海水在阳光照射下波光粼粼，闪烁着钻石般的光芒。天空呈现出清澈的蔚蓝色，点缀着几朵洁白的云朵，营造出晴朗明媚的夏日氛围。画面前景右下角可见粉紫色的小花丛和绿色植物，为整体构图增添了自然生机和色彩层次。整张照片色调明亮清新，紫色头发与白色裙装、蓝色海天形成鲜明而和谐的色彩对比，呈现出一种童话般的浪漫意境，宛如二次元世界与现实海景的完美融合。"
negative_prompt = " "
guidance_scale = 0.00
seed = 43
num_inference_steps = 8
lora_weight = 0.55
save_path = "samples/z-image-t2i-control"
device = set_multi_gpus_devices(ulysses_degree, ring_degree)
config = OmegaConf.load(config_path)
transformer = ZImageControlTransformer2DModel.from_pretrained(
model_name,
subfolder="transformer",
low_cpu_mem_usage=True,
torch_dtype=weight_dtype,
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
if transformer_path is not None:
print(f"From checkpoint: {transformer_path}")
if transformer_path.endswith("safetensors"):
from safetensors.torch import load_file, safe_open
state_dict = load_file(transformer_path)
else:
state_dict = torch.load(transformer_path, map_location="cpu")
state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
m, u = transformer.load_state_dict(state_dict, strict=False)
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
# Get Vae
vae = AutoencoderKL.from_pretrained(
model_name,
subfolder="vae"
).to(weight_dtype)
if vae_path is not None:
print(f"From checkpoint: {vae_path}")
if vae_path.endswith("safetensors"):
from safetensors.torch import load_file, safe_open
state_dict = load_file(vae_path)
else:
state_dict = torch.load(vae_path, map_location="cpu")
state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
m, u = vae.load_state_dict(state_dict, strict=False)
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
# Get tokenizer and text_encoder
tokenizer = AutoTokenizer.from_pretrained(
model_name, subfolder="tokenizer"
)
text_encoder = Qwen3ForCausalLM.from_pretrained(
model_name, subfolder="text_encoder", torch_dtype=weight_dtype,
low_cpu_mem_usage=True,
)
# Get Scheduler
Chosen_Scheduler = scheduler_dict = {
"Flow": FlowMatchEulerDiscreteScheduler,
"Flow_Unipc": FlowUniPCMultistepScheduler,
"Flow_DPM++": FlowDPMSolverMultistepScheduler,
}[sampler_name]
scheduler = Chosen_Scheduler.from_pretrained(
model_name,
subfolder="scheduler"
)
pipeline = ZImageControlPipeline(
vae=vae,
tokenizer=tokenizer,
text_encoder=text_encoder,
transformer=transformer,
scheduler=scheduler,
)
if ulysses_degree > 1 or ring_degree > 1:
from functools import partial
transformer.enable_multi_gpus_inference()
if fsdp_dit:
shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype, module_to_wrapper=list(transformer.layers))
pipeline.transformer = shard_fn(pipeline.transformer)
print("Add FSDP DIT")
if fsdp_text_encoder:
shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype, module_to_wrapper=list(text_encoder.model.layers))
text_encoder = shard_fn(text_encoder)
print("Add FSDP TEXT ENCODER")
if compile_dit:
for i in range(len(pipeline.transformer.transformer_blocks)):
pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
print("Add Compile")
if GPU_memory_mode == "sequential_cpu_offload":
pipeline.enable_sequential_cpu_offload(device=device)
elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
convert_weight_dtype_wrapper(transformer, weight_dtype)
pipeline.enable_model_cpu_offload(device=device)
elif GPU_memory_mode == "model_cpu_offload":
pipeline.enable_model_cpu_offload(device=device)
elif GPU_memory_mode == "model_full_load_and_qfloat8":
convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
convert_weight_dtype_wrapper(transformer, weight_dtype)
pipeline.to(device=device)
else:
pipeline.to(device=device)
generator = torch.Generator(device=device).manual_seed(seed)
if lora_path is not None:
pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
with torch.no_grad():
if inpaint_image is not None:
inpaint_image = get_image_latent(inpaint_image, sample_size=sample_size)[:, :, 0]
else:
inpaint_image = torch.zeros([1, 3, sample_size[0], sample_size[1]])
if mask_image is not None:
mask_image = get_image_latent(mask_image, sample_size=sample_size)[:, :1, 0]
else:
mask_image = torch.ones([1, 1, sample_size[0], sample_size[1]]) * 255
if control_image is not None:
control_image = get_image_latent(control_image, sample_size=sample_size)[:, :, 0]
sample = pipeline(
prompt = prompt,
negative_prompt = negative_prompt,
height = sample_size[0],
width = sample_size[1],
generator = generator,
guidance_scale = guidance_scale,
image = inpaint_image,
mask_image = mask_image,
control_image = control_image,
num_inference_steps = num_inference_steps,
control_context_scale = control_context_scale,
).images
if lora_path is not None:
pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
def save_results():
if not os.path.exists(save_path):
os.makedirs(save_path, exist_ok=True)
index = len([path for path in os.listdir(save_path)]) + 1
prefix = str(index).zfill(8)
video_path = os.path.join(save_path, prefix + ".png")
image = sample[0]
image.save(video_path)
if ulysses_degree * ring_degree > 1:
import torch.distributed as dist
if dist.get_rank() == 0:
save_results()
else:
save_results()
examples/z_image_fun/predict_i2i_tile_2.1.py
+1 -1
View File
@@ -67,7 +67,7 @@ vae_path = None
lora_path = None
# Other params
sample_size = [1728, 992]
sample_size = [1328, 1328]
# Use torch.float16 if GPU does not support torch.bfloat16
# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
examples/z_image_fun/predict_i2i_tile_2.1_lite.py
+242
View File
@@ -0,0 +1,242 @@
import os
import sys
import numpy as np
import torch
from diffusers import FlowMatchEulerDiscreteScheduler
from omegaconf import OmegaConf
from PIL import Image
current_file_path = os.path.abspath(__file__)
project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
for project_root in project_roots:
sys.path.insert(0, project_root) if project_root not in sys.path else None
from videox_fun.dist import set_multi_gpus_devices, shard_model
from videox_fun.models import (AutoencoderKL, AutoTokenizer,
Qwen3ForCausalLM, ZImageControlTransformer2DModel)
from videox_fun.models.cache_utils import get_teacache_coefficients
from videox_fun.pipeline import ZImageControlPipeline
from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
convert_weight_dtype_wrapper)
from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent, get_image_latent, get_image,
get_video_to_video_latent,
save_videos_grid)
# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
# model_full_load means that the entire model will be moved to the GPU.
#
# model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
# and the transformer model has been quantized to float8, which can save more GPU memory.
#
# model_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory.
#
# model_cpu_offload_and_qfloat8 indicates that the entire model will be moved to the CPU after use,
# and the transformer model has been quantized to float8, which can save more GPU memory.
#
# sequential_cpu_offload means that each layer of the model will be moved to the CPU after use,
# resulting in slower speeds but saving a large amount of GPU memory.
GPU_memory_mode = "model_cpu_offload"
# Multi GPUs config
# Please ensure that the product of ulysses_degree and ring_degree equals the number of GPUs used.
# For example, if you are using 8 GPUs, you can set ulysses_degree = 2 and ring_degree = 4.
# If you are using 1 GPU, you can set ulysses_degree = 1 and ring_degree = 1.
ulysses_degree = 1
ring_degree = 1
# Use FSDP to save more GPU memory in multi gpus.
fsdp_dit = False
fsdp_text_encoder = False
# Compile will give a speedup in fixed resolution and need a little GPU memory.
# The compile_dit is not compatible with the fsdp_dit and sequential_cpu_offload.
compile_dit = False
# Config and model path
config_path = "config/z_image/z_image_control_2.1_lite.yaml"
# model path
model_name = "models/Diffusion_Transformer/Z-Image-Turbo"
# Choose the sampler in "Flow", "Flow_Unipc", "Flow_DPM++"
sampler_name = "Flow"
# Load pretrained model if need
transformer_path = "models/Personalized_Model/Z-Image-Turbo-Fun-Controlnet-Tile-2.1-lite-2601-8steps.safetensors"
vae_path = None
lora_path = None
# Other params
sample_size = [1328, 1328]
# Use torch.float16 if GPU does not support torch.bfloat16
# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
weight_dtype = torch.bfloat16
control_image = "asset/low_res.png"
# The inpaint_image and mask_image is useless in tile model, just set them to None.
inpaint_image = None
mask_image = None
control_context_scale = 0.85
# Please use as detailed a prompt as possible to describe the object that needs to be generated.
prompt = "这是一张充满都市气息的户外人物肖像照片。画面中是一位年轻男性，他展现出时尚而自信的形象。人物拥有精心打理的短发发型，两侧修剪得较短，顶部保留一定长度，呈现出流行的Undercut造型。他佩戴着一副时尚的浅色墨镜或透明镜框眼镜，为整体造型增添了潮流感。脸上洋溢着温和友善的笑容，神情放松自然，给人以阳光开朗的印象。他身穿一件经典的牛仔外套，这件单品永不过时，展现出休闲又有型的穿衣风格。牛仔外套的蓝色调与整体氛围十分协调，领口处隐约可见内搭的衣物。照片的背景是典型的城市街景，可以看到模糊的建筑物、街道和行人，营造出繁华都市的氛围。背景经过了恰当的虚化处理，使人物主体更加突出。光线明亮而柔和，可能是白天的自然光，为照片带来清新通透的视觉效果。整张照片构图专业，景深控制得当，完美捕捉了一个现代都市年轻人充满活力和自信的瞬间，展现出积极向上的生活态度。"
negative_prompt = " "
guidance_scale = 0.00
seed = 43
num_inference_steps = 8
lora_weight = 0.55
save_path = "samples/z-image-t2i-control"
device = set_multi_gpus_devices(ulysses_degree, ring_degree)
config = OmegaConf.load(config_path)
transformer = ZImageControlTransformer2DModel.from_pretrained(
model_name,
subfolder="transformer",
low_cpu_mem_usage=True,
torch_dtype=weight_dtype,
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
if transformer_path is not None:
print(f"From checkpoint: {transformer_path}")
if transformer_path.endswith("safetensors"):
from safetensors.torch import load_file, safe_open
state_dict = load_file(transformer_path)
else:
state_dict = torch.load(transformer_path, map_location="cpu")
state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
m, u = transformer.load_state_dict(state_dict, strict=False)
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
# Get Vae
vae = AutoencoderKL.from_pretrained(
model_name,
subfolder="vae"
).to(weight_dtype)
if vae_path is not None:
print(f"From checkpoint: {vae_path}")
if vae_path.endswith("safetensors"):
from safetensors.torch import load_file, safe_open
state_dict = load_file(vae_path)
else:
state_dict = torch.load(vae_path, map_location="cpu")
state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
m, u = vae.load_state_dict(state_dict, strict=False)
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
# Get tokenizer and text_encoder
tokenizer = AutoTokenizer.from_pretrained(
model_name, subfolder="tokenizer"
)
text_encoder = Qwen3ForCausalLM.from_pretrained(
model_name, subfolder="text_encoder", torch_dtype=weight_dtype,
low_cpu_mem_usage=True,
)
# Get Scheduler
Chosen_Scheduler = scheduler_dict = {
"Flow": FlowMatchEulerDiscreteScheduler,
"Flow_Unipc": FlowUniPCMultistepScheduler,
"Flow_DPM++": FlowDPMSolverMultistepScheduler,
}[sampler_name]
scheduler = Chosen_Scheduler.from_pretrained(
model_name,
subfolder="scheduler"
)
pipeline = ZImageControlPipeline(
vae=vae,
tokenizer=tokenizer,
text_encoder=text_encoder,
transformer=transformer,
scheduler=scheduler,
)
if ulysses_degree > 1 or ring_degree > 1:
from functools import partial
transformer.enable_multi_gpus_inference()
if fsdp_dit:
shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype, module_to_wrapper=list(transformer.layers))
pipeline.transformer = shard_fn(pipeline.transformer)
print("Add FSDP DIT")
if fsdp_text_encoder:
shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype, module_to_wrapper=list(text_encoder.model.layers))
text_encoder = shard_fn(text_encoder)
print("Add FSDP TEXT ENCODER")
if compile_dit:
for i in range(len(pipeline.transformer.transformer_blocks)):
pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
print("Add Compile")
if GPU_memory_mode == "sequential_cpu_offload":
pipeline.enable_sequential_cpu_offload(device=device)
elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
convert_weight_dtype_wrapper(transformer, weight_dtype)
pipeline.enable_model_cpu_offload(device=device)
elif GPU_memory_mode == "model_cpu_offload":
pipeline.enable_model_cpu_offload(device=device)
elif GPU_memory_mode == "model_full_load_and_qfloat8":
convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
convert_weight_dtype_wrapper(transformer, weight_dtype)
pipeline.to(device=device)
else:
pipeline.to(device=device)
generator = torch.Generator(device=device).manual_seed(seed)
if lora_path is not None:
pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
with torch.no_grad():
if inpaint_image is not None:
inpaint_image = get_image_latent(inpaint_image, sample_size=sample_size)[:, :, 0]
else:
inpaint_image = torch.zeros([1, 3, sample_size[0], sample_size[1]])
if mask_image is not None:
mask_image = get_image_latent(mask_image, sample_size=sample_size)[:, :1, 0]
else:
mask_image = torch.ones([1, 1, sample_size[0], sample_size[1]]) * 255
if control_image is not None:
control_image = get_image_latent(control_image, sample_size=sample_size)[:, :, 0]
sample = pipeline(
prompt = prompt,
negative_prompt = negative_prompt,
height = sample_size[0],
width = sample_size[1],
generator = generator,
guidance_scale = guidance_scale,
image = inpaint_image,
mask_image = mask_image,
control_image = control_image,
num_inference_steps = num_inference_steps,
control_context_scale = control_context_scale,
).images
if lora_path is not None:
pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
def save_results():
if not os.path.exists(save_path):
os.makedirs(save_path, exist_ok=True)
index = len([path for path in os.listdir(save_path)]) + 1
prefix = str(index).zfill(8)
video_path = os.path.join(save_path, prefix + ".png")
image = sample[0]
image.save(video_path)
if ulysses_degree * ring_degree > 1:
import torch.distributed as dist
if dist.get_rank() == 0:
save_results()
else:
save_results()
examples/z_image_fun/predict_t2i_control_2.1_lite.py
+241
View File
@@ -0,0 +1,241 @@
import os
import sys
import numpy as np
import torch
from diffusers import FlowMatchEulerDiscreteScheduler
from omegaconf import OmegaConf
from PIL import Image
current_file_path = os.path.abspath(__file__)
project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
for project_root in project_roots:
sys.path.insert(0, project_root) if project_root not in sys.path else None
from videox_fun.dist import set_multi_gpus_devices, shard_model
from videox_fun.models import (AutoencoderKL, AutoTokenizer,
Qwen3ForCausalLM, ZImageControlTransformer2DModel)
from videox_fun.models.cache_utils import get_teacache_coefficients
from videox_fun.pipeline import ZImageControlPipeline
from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
convert_weight_dtype_wrapper)
from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent, get_image_latent, get_image,
get_video_to_video_latent,
save_videos_grid)
# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
# model_full_load means that the entire model will be moved to the GPU.
#
# model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
# and the transformer model has been quantized to float8, which can save more GPU memory.
#
# model_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory.
#
# model_cpu_offload_and_qfloat8 indicates that the entire model will be moved to the CPU after use,
# and the transformer model has been quantized to float8, which can save more GPU memory.
#
# sequential_cpu_offload means that each layer of the model will be moved to the CPU after use,
# resulting in slower speeds but saving a large amount of GPU memory.
GPU_memory_mode = "model_cpu_offload"
# Multi GPUs config
# Please ensure that the product of ulysses_degree and ring_degree equals the number of GPUs used.
# For example, if you are using 8 GPUs, you can set ulysses_degree = 2 and ring_degree = 4.
# If you are using 1 GPU, you can set ulysses_degree = 1 and ring_degree = 1.
ulysses_degree = 1
ring_degree = 1
# Use FSDP to save more GPU memory in multi gpus.
fsdp_dit = False
fsdp_text_encoder = False
# Compile will give a speedup in fixed resolution and need a little GPU memory.
# The compile_dit is not compatible with the fsdp_dit and sequential_cpu_offload.
compile_dit = False
# Config and model path
config_path = "config/z_image/z_image_control_2.1_lite.yaml"
# model path
model_name = "models/Diffusion_Transformer/Z-Image-Turbo"
# Choose the sampler in "Flow", "Flow_Unipc", "Flow_DPM++"
sampler_name = "Flow"
# Load pretrained model if need
transformer_path = "models/Personalized_Model/Z-Image-Turbo-Fun-Controlnet-Union-2.1-lite-2601-8steps.safetensors"
vae_path = None
lora_path = None
# Other params
sample_size = [1728, 992]
# Use torch.float16 if GPU does not support torch.bfloat16
# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
weight_dtype = torch.bfloat16
control_image = "asset/pose.jpg"
inpaint_image = None
mask_image = None
control_context_scale = 0.85
# 使用更长的neg prompt如"模糊，突变，变形，失真，画面暗，文本字幕，画面固定，连环画，漫画，线稿，没有主体。"，可以增加稳定性
prompt = "画面中央是一位年轻女孩，她拥有一头令人印象深刻的亮紫色长发，发丝在海风中轻盈飘扬，营造出动感而唯美的效果。她的长发两侧各扎着黑色蝴蝶结发饰，增添了几分可爱与俏皮感。女孩身穿一袭纯白色无袖连衣裙，裙摆轻盈飘逸，与她清新的气质完美契合。她的妆容精致自然，淡粉色的唇妆和温柔的眼神流露出恬静优雅的气质。她单手叉腰，姿态自信从容，目光直视镜头，展现出既甜美又不失个性的魅力。背景是一片开阔的海景，湛蓝的海水在阳光照射下波光粼粼，闪烁着钻石般的光芒。天空呈现出清澈的蔚蓝色，点缀着几朵洁白的云朵，营造出晴朗明媚的夏日氛围。画面前景右下角可见粉紫色的小花丛和绿色植物，为整体构图增添了自然生机和色彩层次。整张照片色调明亮清新，紫色头发与白色裙装、蓝色海天形成鲜明而和谐的色彩对比，呈现出一种童话般的浪漫意境，宛如二次元世界与现实海景的完美融合。"
negative_prompt = " "
guidance_scale = 0.00
seed = 43
num_inference_steps = 8
lora_weight = 0.55
save_path = "samples/z-image-t2i-control"
device = set_multi_gpus_devices(ulysses_degree, ring_degree)
config = OmegaConf.load(config_path)
transformer = ZImageControlTransformer2DModel.from_pretrained(
model_name,
subfolder="transformer",
low_cpu_mem_usage=True,
torch_dtype=weight_dtype,
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
).to(weight_dtype)
if transformer_path is not None:
print(f"From checkpoint: {transformer_path}")
if transformer_path.endswith("safetensors"):
from safetensors.torch import load_file, safe_open
state_dict = load_file(transformer_path)
else:
state_dict = torch.load(transformer_path, map_location="cpu")
state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
m, u = transformer.load_state_dict(state_dict, strict=False)
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
# Get Vae
vae = AutoencoderKL.from_pretrained(
model_name,
subfolder="vae"
).to(weight_dtype)
if vae_path is not None:
print(f"From checkpoint: {vae_path}")
if vae_path.endswith("safetensors"):
from safetensors.torch import load_file, safe_open
state_dict = load_file(vae_path)
else:
state_dict = torch.load(vae_path, map_location="cpu")
state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
m, u = vae.load_state_dict(state_dict, strict=False)
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
# Get tokenizer and text_encoder
tokenizer = AutoTokenizer.from_pretrained(
model_name, subfolder="tokenizer"
)
text_encoder = Qwen3ForCausalLM.from_pretrained(
model_name, subfolder="text_encoder", torch_dtype=weight_dtype,
low_cpu_mem_usage=True,
)
# Get Scheduler
Chosen_Scheduler = scheduler_dict = {
"Flow": FlowMatchEulerDiscreteScheduler,
"Flow_Unipc": FlowUniPCMultistepScheduler,
"Flow_DPM++": FlowDPMSolverMultistepScheduler,
}[sampler_name]
scheduler = Chosen_Scheduler.from_pretrained(
model_name,
subfolder="scheduler"
)
pipeline = ZImageControlPipeline(
vae=vae,
tokenizer=tokenizer,
text_encoder=text_encoder,
transformer=transformer,
scheduler=scheduler,
)
if ulysses_degree > 1 or ring_degree > 1:
from functools import partial
transformer.enable_multi_gpus_inference()
if fsdp_dit:
shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype, module_to_wrapper=list(transformer.layers))
pipeline.transformer = shard_fn(pipeline.transformer)
print("Add FSDP DIT")
if fsdp_text_encoder:
shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype, module_to_wrapper=list(text_encoder.model.layers))
text_encoder = shard_fn(text_encoder)
print("Add FSDP TEXT ENCODER")
if compile_dit:
for i in range(len(pipeline.transformer.transformer_blocks)):
pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i])
print("Add Compile")
if GPU_memory_mode == "sequential_cpu_offload":
pipeline.enable_sequential_cpu_offload(device=device)
elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
convert_weight_dtype_wrapper(transformer, weight_dtype)
pipeline.enable_model_cpu_offload(device=device)
elif GPU_memory_mode == "model_cpu_offload":
pipeline.enable_model_cpu_offload(device=device)
elif GPU_memory_mode == "model_full_load_and_qfloat8":
convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device)
convert_weight_dtype_wrapper(transformer, weight_dtype)
pipeline.to(device=device)
else:
pipeline.to(device=device)
generator = torch.Generator(device=device).manual_seed(seed)
if lora_path is not None:
pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
with torch.no_grad():
if inpaint_image is not None:
inpaint_image = get_image_latent(inpaint_image, sample_size=sample_size)[:, :, 0]
else:
inpaint_image = torch.zeros([1, 3, sample_size[0], sample_size[1]])
if mask_image is not None:
mask_image = get_image_latent(mask_image, sample_size=sample_size)[:, :1, 0]
else:
mask_image = torch.ones([1, 1, sample_size[0], sample_size[1]]) * 255
if control_image is not None:
control_image = get_image_latent(control_image, sample_size=sample_size)[:, :, 0]
sample = pipeline(
prompt = prompt,
negative_prompt = negative_prompt,
height = sample_size[0],
width = sample_size[1],
generator = generator,
guidance_scale = guidance_scale,
image = inpaint_image,
mask_image = mask_image,
control_image = control_image,
num_inference_steps = num_inference_steps,
control_context_scale = control_context_scale,
).images
if lora_path is not None:
pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
def save_results():
if not os.path.exists(save_path):
os.makedirs(save_path, exist_ok=True)
index = len([path for path in os.listdir(save_path)]) + 1
prefix = str(index).zfill(8)
video_path = os.path.join(save_path, prefix + ".png")
image = sample[0]
image.save(video_path)
if ulysses_degree * ring_degree > 1:
import torch.distributed as dist
if dist.get_rank() == 0:
save_results()
else:
save_results()
videox_fun/models/qwenimage_transformer2d.py
+12 -1
View File
@@ -1223,7 +1223,13 @@ class QwenImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Fro
for key in missing_keys:
param_shape = model_state_dict[key].shape
param_dtype = torch_dtype if torch_dtype is not None else model_state_dict[key].dtype
if 'weight' in key:
if "control" in key and key.replace("control_", "transformer_") in filtered_state_dict.keys() and model.state_dict()[key].size() == filtered_state_dict[key.replace("control_", "transformer_")].size():
initialized_dict[key] = filtered_state_dict[key.replace("control_", "transformer_")].clone()
print(f"Initializing missing parameter '{key}' with model.state_dict().")
elif "after_proj" in key or "before_proj" in key:
initialized_dict[key] = torch.zeros(param_shape, dtype=param_dtype)
print(f"Initializing missing parameter '{key}' with zero.")
elif 'weight' in key:
if any(norm_type in key for norm_type in ['norm', 'ln_', 'layer_norm', 'group_norm', 'batch_norm']):
initialized_dict[key] = torch.ones(param_shape, dtype=param_dtype)
elif 'embedding' in key or 'embed' in key:
@@ -1312,6 +1318,11 @@ class QwenImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Fro
tmp_state_dict[key] = state_dict[key]
else:
print(key, "Size don't match, skip")
for key in model.state_dict():
if "control" in key and key.replace("control_", "transformer_") in state_dict.keys() and model.state_dict()[key].size() == state_dict[key.replace("control_", "transformer_")].size():
tmp_state_dict[key] = state_dict[key.replace("control_", "transformer_")].clone()
print(f"Initializing missing parameter '{key}' with model.state_dict().")
state_dict = tmp_state_dict