From d62ae3d2ca0ef6f1a927b62a211b75e63080c2af Mon Sep 17 00:00:00 2001
From: Bubbliiiing <47347516+bubbliiiing@users.noreply.github.com>
Date: Sat, 13 Dec 2025 09:19:16 +0800
Subject: [PATCH] Fix bug in Z image Control Training (#403)

---
 scripts/z_image_fun/train_control.py     | 2 --
 scripts/z_image_fun/train_control_2.0.sh | 1 +
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/scripts/z_image_fun/train_control.py b/scripts/z_image_fun/train_control.py
index 86fde22..484f625 100644
--- a/scripts/z_image_fun/train_control.py
+++ b/scripts/z_image_fun/train_control.py
@@ -86,7 +86,6 @@ from videox_fun.models import (AutoencoderKL, AutoProcessor, AutoTokenizer,
 from videox_fun.pipeline import Flux2Pipeline
 from videox_fun.utils.discrete_sampler import DiscreteSampling
 from videox_fun.utils.utils import get_image_to_video_latent, save_videos_grid
-from videox_fun.utils.utils_yolo import ObjectInstanceDetector
 
 if is_wandb_available():
     import wandb
@@ -842,7 +841,6 @@ def main():
         low_cpu_mem_usage=True,
         transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
     ).to(weight_dtype)
-    yolo_instance = ObjectInstanceDetector(device=accelerator.device)
 
     # Freeze vae and text_encoder and set transformer3d to trainable
     vae.requires_grad_(False)
diff --git a/scripts/z_image_fun/train_control_2.0.sh b/scripts/z_image_fun/train_control_2.0.sh
index f0f1034..86a4240 100644
--- a/scripts/z_image_fun/train_control_2.0.sh
+++ b/scripts/z_image_fun/train_control_2.0.sh
@@ -30,5 +30,6 @@ accelerate launch --mixed_precision="bf16" scripts/z_image_fun/train_control.py
   --max_grad_norm=0.05 \
   --enable_bucket \
   --uniform_sampling \
+  --add_inpaint_info \
   --transformer_path="models/Personalized_Model/Z-Image-Turbo-Fun-Controlnet-Union-2.0.safetensors" \
   --trainable_modules "control"
\ No newline at end of file