Update inference.py

inference.py

@@ -1,6 +1,6 @@
 # inference.py
 import os
-import sys  # 导入 sys 模块
+import sys
 import argparse
 import random
 import json
@@ -13,6 +13,7 @@ import numpy as np
 import requests
 from PIL import Image
 from diffusers import QwenImageEditPipeline
+from huggingface_hub import snapshot_download # <--- 新增导入
 
 # --- 从原脚本保留的辅助函数 ---
 # SYSTEM_PROMPT, polish_prompt, encode_image, api 函数保持不变...
@@ -77,7 +78,6 @@ Please strictly follow the rewriting rules below:
 '''
 
 def polish_prompt(prompt, img):
-    """使用 DashScope API 重写和优化提示词"""
     if not os.environ.get('DASH_API_KEY'):
         print("[警告] 环境变量 DASH_API_KEY 未设置，将跳过提示词重写。")
         return prompt
@@ -98,13 +98,11 @@ def polish_prompt(prompt, img):
     return prompt
 
 def encode_image(pil_image):
-    """将 PIL 图片编码为 base64 字符串"""
     buffered = BytesIO()
     pil_image.save(buffered, format="PNG")
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
 def api(prompt, img_list, model="qwen-vl-max-latest", kwargs={}):
-    """调用 DashScope API"""
     import dashscope
     api_key = os.environ.get('DASH_API_KEY')
     if not api_key:
@@ -120,7 +118,6 @@ def api(prompt, img_list, model="qwen-vl-max-latest", kwargs={}):
         raise Exception(f'Failed to post: {response}')
 
 def load_image(image_path):
-    """从本地路径或URL加载图片"""
     try:
         if image_path.startswith("http://") or image_path.startswith("https://"):
             response = requests.get(image_path)
@@ -134,25 +131,28 @@ def load_image(image_path):
         print(f"   详细信息: {e}")
         return None
 
+# --- 函数修改处 ---
 def prepare_model():
-    """仅下载并缓存模型，不执行推理"""
-    print("正在准备模型... 如果是首次运行，将开始下载模型文件（约10GB）。")
-    print("请耐心等待，下载速度取决于您的网络状况。")
-    dtype = torch.bfloat16
+    """仅下载模型文件到本地缓存，不加载到内存。"""
+    repo_id = "Qwen/Qwen-Image-Edit"
+    print(f"正在准备从 Hugging Face Hub 下载模型 '{repo_id}'...")
+    print("本操作仅下载文件，不会将模型加载到内存或显存中。")
+    print("如果是首次运行，将开始下载模型文件（约7GB），请耐心等待。")
     try:
-        QwenImageEditPipeline.from_pretrained(
-            "Qwen/Qwen-Image-Edit",
-            torch_dtype=dtype,
-            low_cpu_mem_usage=True # 优化内存使用
+        # 使用 snapshot_download 函数只下载文件，并返回其本地路径
+        snapshot_download(
+            repo_id=repo_id,
+            local_dir_use_symlinks=False, # 建议设置为False以提高兼容性
+            resume_download=True # 支持断点续传
         )
-        print("\n✅ 模型文件已成功准备（下载/加载）到本地缓存。")
+        print(f"\n✅ 模型 '{repo_id}' 的文件已成功下载到本地缓存。")
         return True
     except Exception as e:
-        print(f"\n❌ 错误：模型下载或加载失败。请检查网络连接或磁盘空间。")
+        print(f"\n❌ 错误：模型文件下载失败。请检查您的网络连接或仓库名称 '{repo_id}' 是否正确。")
         print(f"   详细信息: {e}")
         return False
 
-# --- 主推理逻辑 ---
+# --- 主推理逻辑 (保持不变) ---
 def main(args):
     """执行模型推理的主函数"""
     output_dir = "output"
@@ -160,7 +160,7 @@ def main(args):
     dtype = torch.bfloat16
     device = "cuda" if torch.cuda.is_available() else "cpu"
     print(f"使用设备: {device}")
-    print("正在加载 Qwen-Image-Edit 模型...")
+    print("正在加载 Qwen-Image-Edit 模型 (从本地缓存)...") # 更新提示
     try:
         pipe = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit", torch_dtype=dtype).to(device)
         print("模型加载完成。")
@@ -193,15 +193,12 @@ def main(args):
     except Exception as e:
         print(f"❌ 推理过程中发生错误: {e}")
 
-# --- 命令行接口 ---
+# --- 命令行接口 (保持不变) ---
 if __name__ == "__main__":
-    # 新增逻辑：检查是否只运行脚本而不带任何参数
     if len(sys.argv) == 1:
         prepare_model()
         print("任务完成，脚本退出。")
-        sys.exit(0) # 正常退出
-
-    # 如果带有参数，则执行原有的推理流程
+        sys.exit(0)
     parser = argparse.ArgumentParser(description="Qwen 图像编辑命令行工具", epilog="如果不提供任何参数，脚本将只下载模型然后退出。")
     parser.add_argument("--prompt",type=str,required=True,help="必须：用于编辑图像的指令。")
     parser.add_argument("--input_image",type=str,required=True,help="必须：输入图片的本地路径或URL链接。")