Update app.py

app.py

@@ -16,8 +16,25 @@ import gradio as gr
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
-from huggingface_hub import snapshot_download
+
+# 延遲導入重量級庫
+diffusers_loaded = False
+def load_diffusers():
+    global DiffusionPipeline, DPMSolverMultistepScheduler, StableDiffusionPipeline, diffusers_loaded
+    if not diffusers_loaded:
+        from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler, StableDiffusionPipeline
+        diffusers_loaded = True
+        return DiffusionPipeline, DPMSolverMultistepScheduler, StableDiffusionPipeline
+    return DiffusionPipeline, DPMSolverMultistepScheduler, StableDiffusionPipeline
+
+huggingface_hub_loaded = False
+def load_huggingface_hub():
+    global snapshot_download, huggingface_hub_loaded
+    if not huggingface_hub_loaded:
+        from huggingface_hub import snapshot_download
+        huggingface_hub_loaded = True
+        return snapshot_download
+    return snapshot_download
 
 HF_TOKEN = os.getenv("HF_TOKEN", "")
 PUBLIC_API_KEY = os.getenv("PUBLIC_API_KEY", "demo-key-123456")
@@ -45,7 +62,7 @@ def get_writable_dir() -> Path:
 PERSIST_DIR = get_writable_dir()
 MODELS_DB_FILE = PERSIST_DIR / "model_library.json"
 HF_CACHE_DIR = Path.home() / ".cache" / "huggingface" / "hub"
-print(f"📁 模型目录: {PERSIST_DIR}")
+print(f"📁 模型目錄: {PERSIST_DIR}")
 
 class ModelLibrary:
     def __init__(self):
@@ -57,6 +74,7 @@ class ModelLibrary:
             try:
                 with open(MODELS_DB_FILE, "r", encoding="utf-8") as f:
                     self.models = json.load(f)
+                print(f"✅ 已載入 {len(self.models)} 個模型記錄")
             except Exception:
                 self.models = {}
         else:
@@ -67,7 +85,7 @@ class ModelLibrary:
             with open(MODELS_DB_FILE, "w", encoding="utf-8") as f:
                 json.dump(self.models, f, indent=2, ensure_ascii=False)
         except Exception as e:
-            print(f"保存失败: {e}")
+            print(f"保存失敗: {e}")
 
     def _key(self, repo_id: str) -> str:
         return repo_id.replace("/", "--").replace(":", "--")
@@ -94,7 +112,7 @@ class ModelLibrary:
 
     def list_display(self) -> List[str]:
         if not self.models:
-            return ["暂无已下载模型"]
+            return ["暫無已下載模型"]
         out = []
         for m in self.models.values():
             emoji = "✅" if m.get("status") == "ready" else "⚠️"
@@ -112,7 +130,7 @@ model_library = ModelLibrary()
 
 class ModelManager:
     def __init__(self):
-        self.loaded: Dict[str, DiffusionPipeline] = {}
+        self.loaded: Dict[str, any] = {}
         self.current: Optional[str] = None
 
     def local_path(self, repo_id: str) -> Path:
@@ -120,7 +138,6 @@ class ModelManager:
 
     def is_downloaded(self, repo_id: str) -> bool:
         lp = self.local_path(repo_id)
-        # Civitai 模型檢查 .safetensors 或 model_index.json
         return (lp / "model_index.json").exists() or any(lp.glob("*.safetensors"))
 
     def size_gb(self, repo_id: str) -> float:
@@ -144,36 +161,22 @@ class ModelManager:
         }
 
     def parse_civitai_url(self, url: str) -> Optional[Dict]:
-        """解析 Civitai 模型 URL"""
         try:
-            # 支援格式：
-            # https://civitai.com/models/12345
-            # https://civitai.com/api/download/models/12345
-            # https://civitai.com/models/12345?modelVersionId=67890
-            
             if "civitai.com" not in url:
                 return None
-            
             parsed = urlparse(url)
-            
-            # 從 URL 提取 model_id 和 version_id
             if "/models/" in url:
                 parts = url.split("/models/")[1].split("/")[0].split("?")[0]
                 model_id = parts
-                
-                # 檢查是否有 modelVersionId
                 query_params = parse_qs(parsed.query)
                 version_id = query_params.get("modelVersionId", [None])[0]
-                
                 return {"model_id": model_id, "version_id": version_id, "url": url}
-            
             return None
         except Exception as e:
             print(f"解析 Civitai URL 失敗: {e}")
             return None
 
     def download_civitai_model(self, url: str) -> tuple:
-        """從 Civitai 下載模型"""
         parsed = self.parse_civitai_url(url)
         if not parsed:
             return False, "❌ 無效的 Civitai URL"
@@ -182,7 +185,6 @@ class ModelManager:
         version_id = parsed["version_id"]
         
         try:
-            # 獲取模型信息
             api_url = f"https://civitai.com/api/v1/models/{model_id}"
             headers = {}
             if CIVITAI_API_KEY:
@@ -194,25 +196,20 @@ class ModelManager:
             model_info = response.json()
             
             model_name = model_info.get("name", f"civitai-{model_id}")
-            
-            # 選擇版本
             versions = model_info.get("modelVersions", [])
             if not versions:
                 return False, "❌ 找不到模型版本"
             
-            # 如果指定了 version_id，使用指定版本；否則使用最新版本
             selected_version = None
             if version_id:
                 selected_version = next((v for v in versions if str(v["id"]) == version_id), None)
             if not selected_version:
                 selected_version = versions[0]
             
-            # 獲取下載 URL
             files = selected_version.get("files", [])
             if not files:
                 return False, "❌ 找不到下載文件"
             
-            # 選擇主文件（通常是 .safetensors）
             download_file = None
             for f in files:
                 if f.get("type") == "Model" or f["name"].endswith(".safetensors"):
@@ -224,9 +221,8 @@ class ModelManager:
             
             download_url = download_file["downloadUrl"]
             file_name = download_file["name"]
-            file_size = download_file.get("sizeKB", 0) / 1024 / 1024  # 轉換為 GB
+            file_size = download_file.get("sizeKB", 0) / 1024 / 1024
             
-            # 創建本地目錄
             repo_id = f"civitai:{model_id}"
             if version_id:
                 repo_id += f":{version_id}"
@@ -234,7 +230,6 @@ class ModelManager:
             local_path = self.local_path(repo_id)
             local_path.mkdir(parents=True, exist_ok=True)
             
-            # 下載文件
             output_file = local_path / file_name
             if output_file.exists():
                 model_library.add(repo_id, model_name, str(local_path), source="civitai")
@@ -243,13 +238,11 @@ class ModelManager:
             print(f"⬇️ 下載 Civitai 模型: {model_name} ({file_size:.2f} GB)")
             start = time.time()
             
-            # 添加 API key 到下載 URL
             if CIVITAI_API_KEY and "?" in download_url:
                 download_url += f"&token={CIVITAI_API_KEY}"
             elif CIVITAI_API_KEY:
                 download_url += f"?token={CIVITAI_API_KEY}"
             
-            # 流式下載
             response = requests.get(download_url, stream=True, timeout=300)
             response.raise_for_status()
             
@@ -261,9 +254,6 @@ class ModelManager:
                     if chunk:
                         f.write(chunk)
                         downloaded += len(chunk)
-                        if total_size > 0 and downloaded % (50 * 1024 * 1024) == 0:  # 每 50MB 打印一次
-                            progress = (downloaded / total_size) * 100
-                            print(f"進度: {progress:.1f}%")
             
             elapsed = time.time() - start
             actual_size = self.size_gb(repo_id)
@@ -278,11 +268,9 @@ class ModelManager:
             return False, f"❌ 下載失敗: {str(e)[:150]}"
 
     def download(self, repo_id: str):
-        # 檢查是否為 Civitai URL
         if "civitai.com" in repo_id:
             return self.download_civitai_model(repo_id)
         
-        # 原有的 Hugging Face 下載邏輯
         if self.is_downloaded(repo_id):
             lp = self.local_path(repo_id)
             model_library.add(repo_id, repo_id.split("/")[-1], str(lp), source="huggingface")
@@ -292,8 +280,9 @@ class ModelManager:
         print(f"⬇️ 下載 HF 模型: {repo_id}")
         
         try:
+            snapshot_download_func = load_huggingface_hub()
             start = time.time()
-            snapshot_download(
+            snapshot_download_func(
                 repo_id=repo_id,
                 local_dir=str(lp),
                 local_dir_use_symlinks=False,
@@ -324,13 +313,11 @@ class ModelManager:
         
         lp = self.local_path(repo_id)
         try:
-            print(f"⚙️ 加載: {lp}")
+            print(f"⚙️ 加載模型: {lp}")
+            DiffusionPipeline, DPMSolverMultistepScheduler, StableDiffusionPipeline = load_diffusers()
             
-            # 檢查是否為 Civitai 單文件模型
             safetensors_files = list(lp.glob("*.safetensors"))
             if safetensors_files and not (lp / "model_index.json").exists():
-                # 單文件 safetensors 加載
-                from diffusers import StableDiffusionPipeline
                 pipe = StableDiffusionPipeline.from_single_file(
                     str(safetensors_files[0]),
                     torch_dtype=torch.float32,
@@ -339,7 +326,6 @@ class ModelManager:
                     low_cpu_mem_usage=True
                 )
             else:
-                # 標準 Diffusers 格式加載
                 pipe = DiffusionPipeline.from_pretrained(
                     str(lp),
                     torch_dtype=torch.float32,
@@ -383,7 +369,6 @@ class ModelManager:
 
 model_manager = ModelManager()
 
-# FastAPI 和 Gradio 部分保持不變...
 app = FastAPI()
 
 class GenerateRequest(BaseModel):
@@ -443,7 +428,7 @@ def ui_generate(prompt, model_display, repo_input, negative, guidance, steps, ap
     if not prompt or len(prompt) < 3:
         return None, "❌ 提示詞太短"
     repo_id = repo_input.strip() if repo_input and len(repo_input.strip()) > 3 else model_library.resolve_repo(model_display)
-    if not repo_id or "暂无" in repo_id:
+    if not repo_id or "暫無" in repo_id:
         return None, "❌ 請選擇或輸入模型"
     ok, msg, pipe = model_manager.load(repo_id)
     if not ok or pipe is None:
@@ -469,35 +454,32 @@ def ui_download(repo_id, api_key):
 def ui_delete(model_display, api_key):
     if api_key != PUBLIC_API_KEY:
         return "❌ API Key 無效", gr.update(choices=model_library.list_display()), model_manager.storage_info()
-    if not model_display or "暂无" in model_display:
+    if not model_display or "暫無" in model_display:
         return "❌ 請先選擇模型", gr.update(choices=model_library.list_display()), model_manager.storage_info()
     repo_id = model_library.resolve_repo(model_display)
     ok, msg = model_manager.delete(repo_id)
     return msg, gr.update(choices=model_library.list_display()), model_manager.storage_info()
 
 def ui_update_api_examples(prompt, model_display, repo_input, negative, guidance, steps):
-    repo_id = repo_input.strip() if repo_input and len(repo_input.strip()) > 3 else (model_library.resolve_repo(model_display) if model_display and "暂无" not in model_display else "prompthero/openjourney-v4")
+    repo_id = repo_input.strip() if repo_input and len(repo_input.strip()) > 3 else (model_library.resolve_repo(model_display) if model_display and "暫無" not in model_display else "prompthero/openjourney-v4")
     return build_api_examples(prompt, repo_id, negative, guidance, steps)
 
-with gr.Blocks(title="AI 圖片生成") as demo:
+with gr.Blocks(title="AI 圖片生成", theme=gr.themes.Soft()) as demo:
     gr.Markdown(f"""
-    # 🖼️ AI 圖片生成（支援 Hugging Face + Civitai）
-    - 模型目錄: `{PERSIST_DIR}`
-    - API Key: `{PUBLIC_API_KEY}`
-    
-    ## 📥 支援模型源
-    - 🤗 **Hugging Face**: `prompthero/openjourney-v4`, `Lykon/dreamshaper-8`
-    - 🎨 **Civitai**: 貼上模型頁面 URL，例如 `https://civitai.com/models/12345`
+    # 🖼️ AI 圖片生成（HF + Civitai）
+    - 📁 模型目錄: `{PERSIST_DIR}`
+    - 🔑 API Key: `{PUBLIC_API_KEY}`
+    - 🤗 **HF**: `prompthero/openjourney-v4` · 🎨 **Civitai**: `https://civitai.com/models/xxx`
     """)
     
     with gr.Tabs():
-        with gr.Tab("🎨 生圖 + API"):
+        with gr.Tab("🎨 生圖"):
             with gr.Row():
                 with gr.Column():
-                    model_dropdown = gr.Dropdown(label="已下載模型 (🤗HF / 🎨Civitai)", choices=model_library.list_display(), value=model_library.list_display()[0] if model_library.list_display() else None)
-                    repo_box = gr.Textbox(label="或輸入模型地址 / Civitai URL", placeholder="prompthero/openjourney-v4 或 https://civitai.com/models/xxx")
-                    prompt_box = gr.Textbox(label="提示詞", lines=3)
-                    negative_box = gr.Textbox(label="負面提示詞", lines=2)
+                    model_dropdown = gr.Dropdown(label="已下載模型", choices=model_library.list_display(), value=model_library.list_display()[0] if model_library.list_display() else None)
+                    repo_box = gr.Textbox(label="或輸入模型地址 / Civitai URL", placeholder="prompthero/openjourney-v4")
+                    prompt_box = gr.Textbox(label="提示詞", lines=3, placeholder="a beautiful sunset...")
+                    negative_box = gr.Textbox(label="負面提示詞", lines=2, placeholder="blurry, low quality...")
                     with gr.Row():
                         guidance_slider = gr.Slider(1, 20, 7.5, step=0.5, label="Guidance")
                         steps_slider = gr.Slider(4, 50, 25, step=1, label="Steps")
@@ -506,20 +488,20 @@ with gr.Blocks(title="AI 圖片生成") as demo:
                 with gr.Column():
                     out_img = gr.Image(label="結果")
                     out_msg = gr.Textbox(label="狀態", lines=5)
-            gr.Markdown("### API 示例")
-            api_json = gr.Code(label="JSON", language="json")
-            api_curl = gr.Code(label="cURL", language="shell")
+            
+            with gr.Accordion("📡 API 示例", open=False):
+                api_json = gr.Code(label="JSON", language="json")
+                api_curl = gr.Code(label="cURL", language="shell")
+            
             gen_btn.click(fn=ui_generate, inputs=[prompt_box, model_dropdown, repo_box, negative_box, guidance_slider, steps_slider, api_key_box], outputs=[out_img, out_msg]).then(fn=lambda: gr.update(choices=model_library.list_display()), outputs=[model_dropdown])
             for comp in [prompt_box, model_dropdown, repo_box, negative_box, guidance_slider, steps_slider]:
                 comp.change(fn=ui_update_api_examples, inputs=[prompt_box, model_dropdown, repo_box, negative_box, guidance_slider, steps_slider], outputs=[api_json, api_curl])
-            demo.load(fn=ui_update_api_examples, inputs=[prompt_box, model_dropdown, repo_box, negative_box, guidance_slider, steps_slider], outputs=[api_json, api_curl])
         
         with gr.Tab("⬇️ 模型管理"):
             with gr.Row():
                 with gr.Column():
                     gr.Markdown("### 下載新模型")
-                    gr.Markdown("支援:\n- 🤗 HF: `prompthero/openjourney-v4`\n- 🎨 Civitai: `https://civitai.com/models/xxx`")
-                    dl_repo = gr.Textbox(label="模型地址 / Civitai URL", placeholder="prompthero/openjourney-v4 或 https://civitai.com/models/xxx")
+                    dl_repo = gr.Textbox(label="模型地址 / URL", placeholder="prompthero/openjourney-v4")
                     dl_key = gr.Textbox(label="API Key", value=PUBLIC_API_KEY, type="password")
                     dl_btn = gr.Button("⬇️ 下載")
                     dl_msg = gr.Textbox(label="結果", lines=4)
@@ -535,9 +517,11 @@ with gr.Blocks(title="AI 圖片生成") as demo:
             del_btn.click(fn=ui_delete, inputs=[lib_dropdown, del_key], outputs=[del_msg, lib_dropdown, storage_box]).then(fn=lambda: gr.update(choices=model_library.list_display()), outputs=[model_dropdown])
             refresh_btn.click(fn=lambda: gr.update(choices=model_library.list_display()), outputs=[lib_dropdown]).then(fn=model_manager.storage_info, outputs=[storage_box]).then(fn=lambda: gr.update(choices=model_library.list_display()), outputs=[model_dropdown])
 
+print("🚀 啟動中...")
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 app = gr.mount_gradio_app(app, demo, path="/")
 
 if __name__ == "__main__":
     import uvicorn
+    print("✅ 準備完成，正在啟動服務器...")
     uvicorn.run(app, host="0.0.0.0", port=7860)