Implement target prefix/suffix functionality in app.py and create_image_caption_json.py for improved image processing. Update UI to reflect changes, enhancing user experience with clearer labeling and file handling. Refactor related functions to accommodate new naming conventions.

app.py

@@ -193,6 +193,24 @@ def _list_checkpoints(out_dir: str, limit: int = 20) -> List[str]:
         return []
 
 
+def _files_to_gallery(files: Any) -> List[str]:
+    items: List[str] = []
+    if not files:
+        return items
+    seq = files if isinstance(files, (list, tuple)) else [files]
+    for f in seq:
+        p = None
+        if isinstance(f, str):
+            p = f
+        elif isinstance(f, dict):
+            p = f.get("path") or f.get("name")
+        else:
+            p = getattr(f, "path", None) or getattr(f, "name", None)
+        if p:
+            items.append(p)
+    return items
+
+
 def _prepare_script(
     dataset_name: str,
     caption: str,
@@ -205,8 +223,8 @@ def _prepare_script(
     override_max_epochs: Optional[int] = None,
     override_save_every: Optional[int] = None,
     override_run_name: Optional[str] = None,
-    main_prefix: Optional[str] = None,
-    main_suffix: Optional[str] = None,
+    target_prefix: Optional[str] = None,
+    target_suffix: Optional[str] = None,
     control_prefixes: Optional[List[Optional[str]]] = None,
     control_suffixes: Optional[List[Optional[str]]] = None,
 ) -> Path:
@@ -300,10 +318,10 @@ def _prepare_script(
 
     # Inject prefix/suffix flags for metadata creation
     extra_lines: List[str] = []
-    if (main_prefix or ""):
-        extra_lines.append(f"  --main_prefix {_bash_quote(main_prefix)} \\")
-    if (main_suffix or ""):
-        extra_lines.append(f"  --main_suffix {_bash_quote(main_suffix)} \\")
+    if (target_prefix or ""):
+        extra_lines.append(f"  --target_prefix {_bash_quote(target_prefix)} \\")
+    if (target_suffix or ""):
+        extra_lines.append(f"  --target_suffix {_bash_quote(target_suffix)} \\")
     for i in range(8):
         pre = control_prefixes[i] if (control_prefixes and i < len(control_prefixes)) else None
         suf = control_suffixes[i] if (control_suffixes and i < len(control_suffixes)) else None
@@ -434,8 +452,8 @@ def run_training(
     output_name: str,
     caption: str,
     image_uploads: Any,
-    main_prefix: str,
-    main_suffix: str,
+    target_prefix: str,
+    target_suffix: str,
     control0_uploads: Any,
     ctrl0_prefix: str,
     ctrl0_suffix: str,
@@ -560,8 +578,8 @@ def run_training(
         override_max_epochs=max_epochs if max_epochs and max_epochs > 0 else None,
         override_save_every=save_every if save_every and save_every > 0 else None,
         override_run_name=output_name.strip(),
-        main_prefix=(main_prefix or ""),
-        main_suffix=(main_suffix or ""),
+        target_prefix=(target_prefix or ""),
+        target_suffix=(target_suffix or ""),
         control_prefixes=[ctrl0_prefix, ctrl1_prefix, ctrl2_prefix, ctrl3_prefix, ctrl4_prefix, ctrl5_prefix, ctrl6_prefix, ctrl7_prefix],
         control_suffixes=[ctrl0_suffix, ctrl1_suffix, ctrl2_suffix, ctrl3_suffix, ctrl4_suffix, ctrl5_suffix, ctrl6_suffix, ctrl7_suffix],
     )
@@ -609,6 +627,15 @@ def build_ui() -> gr.Blocks:
     with gr.Blocks(title="Qwen-Image-Edit: Trainer") as demo:
         gr.Markdown("""
         # Qwen-Image-Edit Trainer
+        このページでは、学習に使う画像をアップロードし、必要ならファイル名の前後にある共通の文字（prefix/suffix）を指定して、
+        自動でデータセットを作成し学習を開始します。難しい操作は不要です。
+
+        画像と名前の対応づけ
+        - prefix/suffix: ファイルの同名判定のため、画像のファイル名から共通の先頭/末尾文字を取り除く指定（例: IMG_ や _v2）
+        - まずメイン画像のファイル名（拡張子なし）から、指定した Target prefix/suffix を取り除いたものを key とします。
+        - 各コントロールは「付加」規則で、期待名 = control_prefix_i + key + control_suffix_i + ".png" を探して対応付けます。
+        - アップロード時に画像は自動で .png に変換して保存します（元のファイル名のベースは維持）。
+        - Control 0 は必須、Control 1〜7 は任意。コントロール画像が1枚だけのときは、すべてのメイン画像に適用します。
         """)
 
         with gr.Row():
@@ -616,90 +643,71 @@ def build_ui() -> gr.Blocks:
             caption = gr.Textbox(label="CAPTION", placeholder="A photo of ...", lines=2)
 
         with gr.Row():
-            images_input = gr.File(label="Upload main images", file_count="multiple", type="filepath", scale=3)
-            try:
-                with gr.Column(scale=1):
-                    main_prefix = gr.Textbox(label="Main prefix", placeholder="e.g., img_")
-                    main_suffix = gr.Textbox(label="Main suffix", placeholder="e.g., _v2")
-            except Exception:
-                # Fallback if Column is unavailable
-                main_prefix = gr.Textbox(label="Main prefix", placeholder="e.g., img_")
-                main_suffix = gr.Textbox(label="Main suffix", placeholder="e.g., _v2")
+            with gr.Column(scale=3):
+                images_input = gr.File(label="Upload target images", file_count="multiple", type="filepath")
+                target_gallery = gr.Gallery(label="Target preview", columns=4, height=200)
+            with gr.Column(scale=1):
+                target_prefix = gr.Textbox(label="Target prefix", placeholder="e.g., IMG_")
+                target_suffix = gr.Textbox(label="Target suffix", placeholder="e.g., _v2")
 
         # control_0 is required and shown outside the accordion
         with gr.Row():
-            ctrl0_files = gr.File(label="Upload control_0 images (required)", file_count="multiple", type="filepath", scale=3)
-            try:
-                with gr.Column(scale=1):
-                    ctrl0_prefix = gr.Textbox(label="control_0 prefix", placeholder="")
-                    ctrl0_suffix = gr.Textbox(label="control_0 suffix", placeholder="")
-            except Exception:
-                ctrl0_prefix = gr.Textbox(label="control_0 prefix", placeholder="")
-                ctrl0_suffix = gr.Textbox(label="control_0 suffix", placeholder="")
+            with gr.Column(scale=3):
+                ctrl0_files = gr.File(label="Upload control_0 images (required)", file_count="multiple", type="filepath")
+                ctrl0_gallery = gr.Gallery(label="control_0 preview", columns=4, height=200)
+            with gr.Column(scale=1):
+                ctrl0_prefix = gr.Textbox(label="control_0 prefix", placeholder="e.g., C0_")
+                ctrl0_suffix = gr.Textbox(label="control_0 suffix", placeholder="e.g., _mask")
 
         # Optional controls start from 1, accordion closed by default
         with gr.Accordion("Optional control images (control_1..control_7)", open=False):
             with gr.Row():
-                ctrl1_files = gr.File(label="Upload control_1 images", file_count="multiple", type="filepath", scale=3)
-                try:
-                    with gr.Column(scale=1):
-                        ctrl1_prefix = gr.Textbox(label="control_1 prefix", placeholder="")
-                        ctrl1_suffix = gr.Textbox(label="control_1 suffix", placeholder="")
-                except Exception:
+                with gr.Column(scale=3):
+                    ctrl1_files = gr.File(label="Upload control_1 images", file_count="multiple", type="filepath")
+                    ctrl1_gallery = gr.Gallery(label="control_1 preview", columns=4, height=200)
+                with gr.Column(scale=1):
                     ctrl1_prefix = gr.Textbox(label="control_1 prefix", placeholder="")
                     ctrl1_suffix = gr.Textbox(label="control_1 suffix", placeholder="")
             with gr.Row():
-                ctrl2_files = gr.File(label="Upload control_2 images", file_count="multiple", type="filepath", scale=3)
-                try:
-                    with gr.Column(scale=1):
-                        ctrl2_prefix = gr.Textbox(label="control_2 prefix", placeholder="")
-                        ctrl2_suffix = gr.Textbox(label="control_2 suffix", placeholder="")
-                except Exception:
+                with gr.Column(scale=3):
+                    ctrl2_files = gr.File(label="Upload control_2 images", file_count="multiple", type="filepath")
+                    ctrl2_gallery = gr.Gallery(label="control_2 preview", columns=4, height=200)
+                with gr.Column(scale=1):
                     ctrl2_prefix = gr.Textbox(label="control_2 prefix", placeholder="")
                     ctrl2_suffix = gr.Textbox(label="control_2 suffix", placeholder="")
             with gr.Row():
-                ctrl3_files = gr.File(label="Upload control_3 images", file_count="multiple", type="filepath", scale=3)
-                try:
-                    with gr.Column(scale=1):
-                        ctrl3_prefix = gr.Textbox(label="control_3 prefix", placeholder="")
-                        ctrl3_suffix = gr.Textbox(label="control_3 suffix", placeholder="")
-                except Exception:
+                with gr.Column(scale=3):
+                    ctrl3_files = gr.File(label="Upload control_3 images", file_count="multiple", type="filepath")
+                    ctrl3_gallery = gr.Gallery(label="control_3 preview", columns=4, height=200)
+                with gr.Column(scale=1):
                     ctrl3_prefix = gr.Textbox(label="control_3 prefix", placeholder="")
                     ctrl3_suffix = gr.Textbox(label="control_3 suffix", placeholder="")
             with gr.Row():
-                ctrl4_files = gr.File(label="Upload control_4 images", file_count="multiple", type="filepath", scale=3)
-                try:
-                    with gr.Column(scale=1):
-                        ctrl4_prefix = gr.Textbox(label="control_4 prefix", placeholder="")
-                        ctrl4_suffix = gr.Textbox(label="control_4 suffix", placeholder="")
-                except Exception:
+                with gr.Column(scale=3):
+                    ctrl4_files = gr.File(label="Upload control_4 images", file_count="multiple", type="filepath")
+                    ctrl4_gallery = gr.Gallery(label="control_4 preview", columns=4, height=200)
+                with gr.Column(scale=1):
                     ctrl4_prefix = gr.Textbox(label="control_4 prefix", placeholder="")
                     ctrl4_suffix = gr.Textbox(label="control_4 suffix", placeholder="")
             with gr.Row():
-                ctrl5_files = gr.File(label="Upload control_5 images", file_count="multiple", type="filepath", scale=3)
-                try:
-                    with gr.Column(scale=1):
-                        ctrl5_prefix = gr.Textbox(label="control_5 prefix", placeholder="")
-                        ctrl5_suffix = gr.Textbox(label="control_5 suffix", placeholder="")
-                except Exception:
+                with gr.Column(scale=3):
+                    ctrl5_files = gr.File(label="Upload control_5 images", file_count="multiple", type="filepath")
+                    ctrl5_gallery = gr.Gallery(label="control_5 preview", columns=4, height=200)
+                with gr.Column(scale=1):
                     ctrl5_prefix = gr.Textbox(label="control_5 prefix", placeholder="")
                     ctrl5_suffix = gr.Textbox(label="control_5 suffix", placeholder="")
             with gr.Row():
-                ctrl6_files = gr.File(label="Upload control_6 images", file_count="multiple", type="filepath", scale=3)
-                try:
-                    with gr.Column(scale=1):
-                        ctrl6_prefix = gr.Textbox(label="control_6 prefix", placeholder="")
-                        ctrl6_suffix = gr.Textbox(label="control_6 suffix", placeholder="")
-                except Exception:
+                with gr.Column(scale=3):
+                    ctrl6_files = gr.File(label="Upload control_6 images", file_count="multiple", type="filepath")
+                    ctrl6_gallery = gr.Gallery(label="control_6 preview", columns=4, height=200)
+                with gr.Column(scale=1):
                     ctrl6_prefix = gr.Textbox(label="control_6 prefix", placeholder="")
                     ctrl6_suffix = gr.Textbox(label="control_6 suffix", placeholder="")
             with gr.Row():
-                ctrl7_files = gr.File(label="Upload control_7 images", file_count="multiple", type="filepath", scale=3)
-                try:
-                    with gr.Column(scale=1):
-                        ctrl7_prefix = gr.Textbox(label="control_7 prefix", placeholder="")
-                        ctrl7_suffix = gr.Textbox(label="control_7 suffix", placeholder="")
-                except Exception:
+                with gr.Column(scale=3):
+                    ctrl7_files = gr.File(label="Upload control_7 images", file_count="multiple", type="filepath")
+                    ctrl7_gallery = gr.Gallery(label="control_7 preview", columns=4, height=200)
+                with gr.Column(scale=1):
                     ctrl7_prefix = gr.Textbox(label="control_7 prefix", placeholder="")
                     ctrl7_suffix = gr.Textbox(label="control_7 suffix", placeholder="")
 
@@ -708,16 +716,26 @@ def build_ui() -> gr.Blocks:
         run_btn = gr.Button("Start Training", variant="primary")
         logs = gr.Textbox(label="Logs", lines=20)
         ckpt_files = gr.Files(label="Checkpoints (live)", interactive=False)
-        lora_file = gr.File(label="Download LoRA (latest)", interactive=False)
 
         with gr.Row():
             max_epochs = gr.Number(label="Max epochs (this run)", value=10, precision=0)
             save_every = gr.Number(label="Save every N epochs", value=5, precision=0)
 
+        # Wire previews
+        images_input.change(fn=_files_to_gallery, inputs=images_input, outputs=target_gallery)
+        ctrl0_files.change(fn=_files_to_gallery, inputs=ctrl0_files, outputs=ctrl0_gallery)
+        ctrl1_files.change(fn=_files_to_gallery, inputs=ctrl1_files, outputs=ctrl1_gallery)
+        ctrl2_files.change(fn=_files_to_gallery, inputs=ctrl2_files, outputs=ctrl2_gallery)
+        ctrl3_files.change(fn=_files_to_gallery, inputs=ctrl3_files, outputs=ctrl3_gallery)
+        ctrl4_files.change(fn=_files_to_gallery, inputs=ctrl4_files, outputs=ctrl4_gallery)
+        ctrl5_files.change(fn=_files_to_gallery, inputs=ctrl5_files, outputs=ctrl5_gallery)
+        ctrl6_files.change(fn=_files_to_gallery, inputs=ctrl6_files, outputs=ctrl6_gallery)
+        ctrl7_files.change(fn=_files_to_gallery, inputs=ctrl7_files, outputs=ctrl7_gallery)
+
         run_btn.click(
             fn=run_training,
             inputs=[
-                output_name, caption, images_input, main_prefix, main_suffix,
+                output_name, caption, images_input, target_prefix, target_suffix,
                 ctrl0_files, ctrl0_prefix, ctrl0_suffix,
                 ctrl1_files, ctrl1_prefix, ctrl1_suffix,
                 ctrl2_files, ctrl2_prefix, ctrl2_suffix,
@@ -728,7 +746,7 @@ def build_ui() -> gr.Blocks:
                 ctrl7_files, ctrl7_prefix, ctrl7_suffix,
                 max_epochs, save_every,
             ],
-            outputs=[logs, ckpt_files, lora_file],
+            outputs=[logs, ckpt_files],
         )
 
     return demo
@@ -783,3 +801,4 @@ if __name__ == "__main__":
     except TypeError:
         # Older gradio without allowed_paths
         ui.launch(server_name="0.0.0.0")
+

create_image_caption_json.py

@@ -87,8 +87,8 @@ def create_image_caption_json_unified(
     caption: str,
     output_json: str,
     *,
-    main_prefix: str = "",
-    main_suffix: str = "",
+    target_prefix: str = "",
+    target_suffix: str = "",
     control_prefixes: Optional[List[Optional[str]]] = None,
     control_suffixes: Optional[List[Optional[str]]] = None,
     allow_single: bool = True,
@@ -105,7 +105,7 @@ def create_image_caption_json_unified(
     filenames = _list_images(input_folder)
 
     # If no prefixes/suffixes provided, keep strict validation behavior
-    use_name_matching = bool(main_prefix or main_suffix)
+    use_name_matching = bool(target_prefix or target_suffix)
     if control_prefixes:
         use_name_matching = True
 
@@ -124,8 +124,8 @@ def create_image_caption_json_unified(
                 f.write(json.dumps(entry, ensure_ascii=False) + '\n')
                 count += 1
         else:
-            # Build main keys
-            main_keys = [_normalize(fn, main_prefix, main_suffix) for fn in filenames]
+            # Build target keys
+            target_keys = [_normalize(fn, target_prefix, target_suffix) for fn in filenames]
 
             # Write matched entries by constructing expected control names using add-rule
             for idx, fname in enumerate(filenames):
@@ -134,7 +134,7 @@ def create_image_caption_json_unified(
                     "caption": caption,
                 }
                 present_count = 0
-                mkey = main_keys[idx]
+                mkey = target_keys[idx]
                 for i, cdir in enumerate(control_dirs):
                     if not cdir:
                         continue
@@ -168,8 +168,8 @@ if __name__ == "__main__":
     parser.add_argument('-o', '--output-json', default='metadata.jsonl', help='出力JSONLパス（既定: metadata.jsonl）')
     parser.add_argument('--image-dir', default='/workspace/data/image', help='image_pathの親ディレクトリパス（JSON出力用）')
 
-    parser.add_argument('--main_prefix', default=os.environ.get('MAIN_PREFIX', ''), help='メイン画像名から削除するプレフィックス')
-    parser.add_argument('--main_suffix', default=os.environ.get('MAIN_SUFFIX', ''), help='メイン画像名から削除するサフィックス')
+    parser.add_argument('--target_prefix', default=os.environ.get('TARGET_PREFIX', ''), help='ターゲット画像名から削除するプレフィックス')
+    parser.add_argument('--target_suffix', default=os.environ.get('TARGET_SUFFIX', ''), help='ターゲット画像名から削除するサフィックス')
 
     # 最大 control_dir_0..7 まで受け付け
     for i in range(8):
@@ -211,8 +211,8 @@ if __name__ == "__main__":
         control_dirs=control_dirs,
         caption=args.caption,
         output_json=args.output_json,
-        main_prefix=args.main_prefix,
-        main_suffix=args.main_suffix,
+        target_prefix=args.target_prefix,
+        target_suffix=args.target_suffix,
         control_prefixes=control_prefixes,
         control_suffixes=control_suffixes,
         allow_single=args.allow_single,