Create colorflow_cli.py

colorflow_cli.py

@@ -0,0 +1,77 @@
+'''
+python colorflow_cli.py \
+  --input_image ./input.jpg \
+  --reference_images ./ref1.jpg ./ref2.jpg \
+  --output_dir ./results \
+  --input_style Sketch \
+  --resolution 640x640 \
+  --seed 123 \
+  --num_inference_steps 20
+'''
+
+# colorflow_cli.py
+from app_func import *
+import argparse
+import torch
+from PIL import Image
+import os
+import logging
+
+# 原文件中的必要导入和函数定义（需保留原文件中的核心逻辑）
+# ... [保留原文件中的模型加载、extract_line_image、colorize_image等函数] ...
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="ColorFlow命令行图像上色工具")
+    parser.add_argument("--input_image", type=str, required=True, help="输入图像路径")
+    parser.add_argument("--reference_images", type=str, nargs='+', required=True, help="参考图像路径列表")
+    parser.add_argument("--output_dir", type=str, default="./output", help="输出目录")
+    parser.add_argument("--input_style", type=str, default="GrayImage(ScreenStyle)", 
+                        choices=["GrayImage(ScreenStyle)", "Sketch"], help="输入样式类型")
+    parser.add_argument("--resolution", type=str, default="640x640", 
+                        choices=["640x640", "512x800", "800x512"], help="分辨率设置")
+    parser.add_argument("--seed", type=int, default=0, help="随机种子")
+    parser.add_argument("--num_inference_steps", type=int, default=10, help="推理步数")
+    return parser.parse_args()
+
+def save_image(image: Image.Image, path: str, format: str = "PNG") -> None:
+    """安全保存图像并处理异常"""
+    try:
+        image.save(path, format=format)
+        logging.info(f"成功保存图像至: {path}")
+    except Exception as e:
+        logging.error(f"保存图像失败: {str(e)}")
+        raise
+
+def main():
+    args = parse_args()
+    os.makedirs(args.output_dir, exist_ok=True)
+    
+    # 初始化模型
+    global cur_input_style, pipeline, MultiResNetModel
+    cur_input_style = None
+    load_ckpt(args.input_style)
+    
+    # 预处理输入图像
+    input_img = Image.open(args.input_image).convert("RGB")
+    input_context, extracted_line, _ = extract_line_image(input_img, args.input_style, args.resolution)
+    
+    # 执行颜色化并获取全部结果
+    high_res_img, up_img, raw_output, preprocessed_bw = colorize_image(
+        VAE_input=extracted_line,
+        input_context=input_context,
+        reference_images=args.reference_images,
+        resolution=args.resolution,
+        seed=args.seed,
+        input_style=args.input_style,
+        num_inference_steps=args.num_inference_steps
+    )
+    
+    # 保存所有结果
+    save_image(high_res_img, os.path.join(args.output_dir, "colorized_result.png"))
+    save_image(up_img, os.path.join(args.output_dir, "upsampled_intermediate.png"))
+    save_image(raw_output, os.path.join(args.output_dir, "raw_generated_output.png"))
+    save_image(preprocessed_bw, os.path.join(args.output_dir, "preprocessed_bw.png"))
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+    main()