22

README.md

@@ -12,6 +12,8 @@ short_description: Supports Canny, HED, Depth, Pose and MLSD
 models:
     - Tongyi-MAI/Z-Image-Turbo
     - alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union
+commands:
+    - echo "Starting Space..."
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,10 +1,6 @@
 import gradio as gr
 import numpy as np
-import random
-import json
-import spaces
-import torch
-from diffusers import DiffusionPipeline
+import random, json, spaces, torch
 from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler
 from videox_fun.pipeline import ZImageControlPipeline
 from videox_fun.models import ZImageControlTransformer2DModel
@@ -12,6 +8,7 @@ from transformers import AutoTokenizer, Qwen3ForCausalLM
 from diffusers import AutoencoderKL
 from utils.image_utils import get_image_latent, scale_image
 from utils.prompt_utils import polish_prompt
+from controlnet_aux import HEDdetector, MLSDdetector, OpenposeDetector, CannyDetector, MidasDetector
 # from videox_fun.utils.utils import get_image_latent
 
 
@@ -61,7 +58,7 @@ tokenizer = AutoTokenizer.from_pretrained(
 )
 text_encoder = Qwen3ForCausalLM.from_pretrained(
     MODEL_LOCAL, subfolder="text_encoder", torch_dtype=weight_dtype,
-    low_cpu_mem_usage=True,
+    low_cpu_mem_usage=False,
 )
 scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3)
 pipe = ZImageControlPipeline(
@@ -79,9 +76,22 @@ pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
 spaces.aoti_blocks_load(pipe.transformer.layers,
                         "zerogpu-aoti/Z-Image", variant="fa3")
 
-def prepare(prompt, input_image):
+def prepare(prompt, input_image, control_mode='Canny'):
     polished_prompt = polish_prompt(prompt)
-    return polished_prompt
+
+    if control_mode == 'HED':
+        processor = HEDdetector.from_pretrained("lllyasviel/Annotators")
+    if control_mode =='Midas':
+        processor = MidasDetector.from_pretrained("lllyasviel/Annotators")
+    if control_mode =='MLSD':
+        processor = MLSDdetector.from_pretrained("lllyasviel/Annotators")
+    if control_mode =='Pose':
+        processor = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
+    else:
+        processor = CannyDetector()
+    control_image = processor(input_image)
+    
+    return polished_prompt, control_image
 
 @spaces.GPU
 def inference(
@@ -141,7 +151,7 @@ with open('static/data.json', 'r') as file:
     data = json.load(file)
 examples = data['examples']
 
-with gr.Blocks() as demo:
+with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         with gr.Column():
             gr.HTML(read_file("static/header.html"))
@@ -244,4 +254,4 @@ with gr.Blocks() as demo:
     # )
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True, css=css)
+    demo.launch(mcp_server=True)

utils/prompt_utils.py

@@ -16,7 +16,6 @@ def polish_prompt(original_prompt):
     # messages = []
     client = InferenceClient()
     
-    
     try:
         completion = client.chat.completions.create(
         rovider="cerebras",