22

app.py

@@ -6,9 +6,7 @@ from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler
 from transformers import AutoTokenizer, Qwen3ForCausalLM
 from safetensors.torch import load_file
 from diffusers import AutoencoderKL
-from utils.image_utils import get_image_latent, rescale_image
-from utils.prompt_utils import polish_prompt
-from utils import repo_utils
+from utils import repo_utils, prompt_utils, image_utils
 from controlnet_aux.processor import Processor
 from omegaconf import OmegaConf
 
@@ -18,6 +16,7 @@ repo_utils.move_folder("repos/VideoX-Fun/videox_fun", "videox_fun")
 repo_utils.move_folder("repos/VideoX-Fun/config", "config")
 from videox_fun.pipeline import ZImageControlPipeline
 from videox_fun.models import ZImageControlTransformer2DModel
+from videox_fun.utils.utils import get_image_latent
 
 #clone models
 repo_utils.clone_repo_if_not_exists("https://huggingface.co/Tongyi-MAI/Z-Image-Turbo", "models")
@@ -80,13 +79,9 @@ pipe = ZImageControlPipeline(
 pipe.to("cuda", torch.bfloat16)
 print("pipe ready.")
 
-# ======== AoTI compilation + FA3 ========
-# pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
-# spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
-
 def prepare(prompt, is_polish_prompt):
     if not is_polish_prompt: return prompt, False
-    polished_prompt = polish_prompt(prompt)
+    polished_prompt =  prompt_utils.polish_prompt(prompt)
     return polished_prompt, True
 
 @spaces.GPU
@@ -128,7 +123,7 @@ def inference(
     processor = Processor(processor_id)
 
     # Width must be divisible by 16
-    control_image, width, height = rescale_image(input_image, image_scale, 16)
+    control_image, width, height = image_utils.rescale_image(input_image, image_scale, 16)
     control_image = control_image.resize((1024, 1024))
 
     print("DEBUG: control_image_torch")
@@ -192,11 +187,19 @@ with gr.Blocks(css=css) as demo:
             gr.HTML(read_file("static/header.html"))
         with gr.Row():
             with gr.Column():
+                mask_image = gr.ImagePaint(
+                    height=290, 
+                    sources=['upload', 'clipboard'], 
+                    image_mode='RGB', 
+                    type="pil", label="Mask Image"
+                )
+
                 input_image = gr.Image(
-                    height=290, sources=['upload', 'clipboard'], 
+                    height=290, 
+                    sources=['upload', 'clipboard'], 
                     image_mode='RGB', 
-                    # elem_id="image_upload", 
-                    type="pil", label="Upload")
+                    type="pil", label="Upload"
+                )
                 
                 prompt = gr.Textbox(
                     label="Prompt",
@@ -273,6 +276,7 @@ with gr.Blocks(css=css) as demo:
         gr.Examples(examples=examples, inputs=[input_image, prompt, control_mode])
         gr.Markdown(read_file("static/footer.md"))
 
+    mask_image.upload(fn=lambda x: x, inputs=[mask_image], outputs=[input_image])
     run_button.click(
         fn=prepare,
         inputs=[prompt, is_polish_prompt],