Update app.py

app.py

@@ -4,7 +4,11 @@ import gc
 import torch
 import requests
 import gradio as gr
-from diffusers import AutoPipelineForText2Image, DPMSolverMultistepScheduler
+from diffusers import AutoPipelineForText2Image, DPMSolverMultistepScheduler, LCMScheduler
+
+# ── 0. CPU 核心效能最佳化 ──────────────────────────────────────────
+# 限制 PyTorch 只使用 2 個執行緒，完美對應 HF 免費空間的 2 vCPU，避免過度切換造成卡頓
+torch.set_num_threads(2)
 
 # ── 1. 設定與全域變數 ──────────────────────────────────────────────
 MODEL_CACHE_DIR = "./custom_models"
@@ -14,17 +18,19 @@ os.makedirs(LORA_CACHE_DIR, exist_ok=True)
 
 pipe = None
 current_model_path = ""
-active_loras = {}  
+is_current_model_sdxl = False
+active_loras = {}  # 存放使用者自訂的 LoRA {"name": scale}
 
 PRESET_MODELS = {
+    "BK-SDM-Tiny (極速輕量 1.5)": "nota-ai/bk-sdm-tiny",
     "Stable Diffusion v1.5 (通用)": "runwayml/stable-diffusion-v1-5",
-    "BK-SDM-Tiny (極速)": "nota-ai/bk-sdm-tiny",
-    "Dreamlike Anime 1.0": "dreamlike-art/dreamlike-anime-1.0",
+    "Dreamlike Anime 1.0 (動漫)": "dreamlike-art/dreamlike-anime-1.0",
 }
 
 # ── 2. 核心邏輯函式 ───────────────────────────────────────────────
 
 def download_file(url, folder, progress, token=""):
+    """支援進度條、Civitai API Token 與防呆檢查的下載器"""
     try:
         headers = {}
         if token and token.strip():
@@ -58,9 +64,9 @@ def download_file(url, folder, progress, token=""):
                 f.write(data)
                 downloaded += len(data)
                 if total_size > 0:
-                    progress(downloaded / total_size, desc=f"下載 {fname}: {downloaded/1024/1024:.1f}MB")
+                    progress(downloaded / total_size, desc=f"下載 {fname[:15]}: {downloaded/1024/1024:.1f}MB")
                 else:
-                    progress(None, desc=f"下載 {fname}...")
+                    progress(None, desc=f"下載 {fname[:15]}...")
                     
         if os.path.exists(filepath) and os.path.getsize(filepath) < 1024 * 100:
             os.remove(filepath)
@@ -72,17 +78,19 @@ def download_file(url, folder, progress, token=""):
 
 
 def load_pipeline(model_source, is_local_file=False):
-    global pipe, current_model_path, active_loras
+    """負責實際載入主模型，並預先準備好 LCM 加速元件"""
+    global pipe, current_model_path, is_current_model_sdxl, active_loras
 
     if model_source == current_model_path and pipe is not None:
         return f"✅ 已載入: {model_source}"
 
+    # 釋放舊模型
     pipe = None
     active_loras = {}
     gc.collect()
 
     try:
-        # 【關鍵修改】使用 AutoPipelineForText2Image，讓它自動判斷是 SD1.5 還是 SDXL
+        # 1. 載入主模型 (自動判斷 SD 1.5 或 SDXL)
         if is_local_file:
             p = AutoPipelineForText2Image.from_single_file(
                 model_source, torch_dtype=torch.float32, 
@@ -94,18 +102,22 @@ def load_pipeline(model_source, is_local_file=False):
                 safety_checker=None, requires_safety_checker=False
             )
 
-        # 設定 Scheduler 並優化 CPU 記憶體
-        p.scheduler = DPMSolverMultistepScheduler.from_config(p.scheduler.config)
         p.to("cpu")
-        p.enable_attention_slicing()
+        p.enable_attention_slicing() # 省記憶體關鍵
+        
+        # 判斷架構以決定使用哪種 LCM-LoRA
+        is_sdxl = "SDXL" in p.__class__.__name__
+        lcm_lora_id = "latent-consistency/lcm-lora-sdxl" if is_sdxl else "latent-consistency/lcm-lora-sdv1-5"
+        
+        # 2. 預先下載並掛載 LCM-LoRA (命名為 "lcm" 以便後續動態開關)
+        p.load_lora_weights(lcm_lora_id, adapter_name="lcm")
+        p.disable_lora() # 預設先關閉，由生成時決定是否開啟
 
         pipe = p
         current_model_path = model_source
-        
-        is_sdxl = "SDXL" in p.__class__.__name__
         model_type_str = "SDXL" if is_sdxl else "SD 1.5"
         
-        return f"✅ 成功載入主模型 ({model_type_str}): {os.path.basename(model_source) if is_local_file else model_source}"
+        return f"✅ 成功載入 ({model_type_str}): {os.path.basename(model_source) if is_local_file else model_source}"
     except Exception as e:
         if is_local_file and os.path.exists(model_source):
             os.remove(model_source)
@@ -113,7 +125,7 @@ def load_pipeline(model_source, is_local_file=False):
 
 
 def load_pipeline_generator(source, is_local):
-    yield "⏳ 載入模型中... (SDXL 在 CPU 上極慢且易崩潰，請耐心等候)"
+    yield "⏳ 載入模型中... (包含下載 LCM 加速模組，請稍候)"
     result = load_pipeline(source, is_local)
     yield result
 
@@ -136,7 +148,7 @@ def handle_civitai_model(url, token, progress=gr.Progress()):
 
 
 def update_lora_list():
-    if not active_loras: return "無啟用 LoRA"
+    if not active_loras: return "無"
     return "\n".join([f"- {k}: {v}" for k, v in active_loras.items()])
 
 
@@ -150,14 +162,11 @@ def add_lora(url, scale, token, progress=gr.Progress()):
         path, fname = download_file(url, LORA_CACHE_DIR, progress, token)
         adapter_name = fname.replace(".", "_")
         
+        # 載入自訂 LoRA
         pipe.load_lora_weights(path, adapter_name=adapter_name)
         active_loras[adapter_name] = float(scale)
         
-        adapters = list(active_loras.keys())
-        weights = list(active_loras.values())
-        pipe.set_adapters(adapters, adapter_weights=weights)
-        
-        return f"✅ 已加入 LoRA: {fname} (權重 {scale})", update_lora_list()
+        return f"✅ 已加入: {fname}", update_lora_list()
     except Exception as e:
         if path and os.path.exists(path):
             os.remove(path)
@@ -165,108 +174,128 @@ def add_lora(url, scale, token, progress=gr.Progress()):
 
 
 def clear_loras():
-    global pipe, active_loras
-    if pipe is None: 
-        return "⚠️ 無模型"
-    try:
-        pipe.unload_lora_weights()
-        active_loras = {}
-        return "🗑️ 已移除所有 LoRA"
-    except Exception as e:
-        return f"❌ 移除失敗: {str(e)}"
+    global active_loras
+    if pipe is None: return "⚠️ 無模型"
+    # 注意：我們不 unload "lcm"，只清除使用者的 active_loras 清單
+    active_loras = {}
+    return "🗑️ 已移除所有自訂 LoRA"
 
 
-def generate_image(prompt, neg, steps, cfg, seed, width, height):
+def generate_image(prompt, neg, steps, cfg, seed, width, height, use_lcm):
+    """執行圖片生成 (包含動態 Scheduler 與 Adapter 切換)"""
     if pipe is None: 
         raise gr.Error("請先載入模型！")
 
+    start_time = time.time()
+
     if seed == -1: 
         seed = int(time.time() % (2**32))
     generator = torch.Generator("cpu").manual_seed(seed)
     
-    # 執行推論
+    # ── 動態切換 LCM 加速與自訂 LoRA ──
+    adapters_to_use = []
+    weights_to_use = []
+
+    if use_lcm:
+        # 切換為 LCM 排程器
+        pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+        adapters_to_use.append("lcm")
+        weights_to_use.append(1.0)
+    else:
+        # 切換為一般高畫質排程器
+        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+
+    # 加入使用者自訂的 LoRA
+    for k, v in active_loras.items():
+        adapters_to_use.append(k)
+        weights_to_use.append(v)
+
+    # 啟用/禁用 Adapters
+    if len(adapters_to_use) > 0:
+        pipe.enable_lora()
+        pipe.set_adapters(adapters_to_use, adapter_weights=weights_to_use)
+    else:
+        pipe.disable_lora()
+
+    # ── 執行生成 ──
     image = pipe(
         prompt=prompt, 
-        negative_prompt=neg, 
+        negative_prompt=neg if not use_lcm else None, # LCM 建議忽略反向提示詞
         num_inference_steps=int(steps), 
-        guidance_scale=cfg, 
+        guidance_scale=float(cfg), 
         width=int(width),
         height=int(height),
         generator=generator
     ).images[0]
     
-    return image, seed
+    cost_time = time.time() - start_time
+    status = f"✅ 完成 | 耗時: {cost_time:.1f} 秒 | Seed: {seed}"
+    return image, status
 
 
 # ── 3. Gradio UI 介面設計 ──────────────────────────────────────────
 
-with gr.Blocks(title="SD/SDXL CPU + Civitai + LoRA") as demo:
-    gr.Markdown("# 🎨 SD/SDXL (CPU) - 支援 Civitai 主模型與 LoRA")
-    gr.Markdown("> ⚠️ 警告：CPU 環境載入 SDXL 模型極度耗時且容易因 16GB 記憶體不足而崩潰，強烈建議使用 SD 1.5 模型。")
+with gr.Blocks(title="Turbo CPU Stable Diffusion") as demo:
+    gr.Markdown("# ⚡ Turbo CPU Stable Diffusion (含 LCM 極速架構)")
+    gr.Markdown("完美適配 HuggingFace 免費 CPU。開啟 LCM 模式可將生成時間從 3 分鐘縮短至 15 秒以內！")
     
     with gr.Row():
+        # ── 左側：模型與管理 ──
         with gr.Column(scale=1):
-            gr.Markdown("### 🔑 Civitai 授權")
-            civit_token = gr.Textbox(
-                label="Civitai API Token (選填)", 
-                placeholder="若需下載 R18 或限定模型，請貼上你的 Token",
-                type="password"
-            )
+            with gr.Accordion("🔑 Civitai 授權 (選填)", open=False):
+                civit_token = gr.Textbox(label="API Token", placeholder="下載限定模型用", type="password")
             
-            gr.Markdown("### 1. 主模型")
+            gr.Markdown("### 1. 選擇主模型")
             with gr.Tabs():
-                with gr.TabItem("📦 預設"):
+                with gr.TabItem("📦 預設 (推薦)"):
                     preset_dd = gr.Dropdown(list(PRESET_MODELS.keys()), label="選擇模型", value=list(PRESET_MODELS.keys())[0])
-                    load_preset_btn = gr.Button("載入預設模型")
-                with gr.TabItem("🌐 Civitai URL"):
-                    civit_ckpt_url = gr.Textbox(label="Checkpoint 下載網址", placeholder="輸入 Civitai 直連下載網址...")
+                    load_preset_btn = gr.Button("載入預設模型", variant="primary")
+                with gr.TabItem("🌐 Civitai 連結"):
+                    civit_ckpt_url = gr.Textbox(label="Checkpoint 網址", placeholder="https://civitai.com/api/...")
                     load_civit_btn = gr.Button("下載並載入")
             
-            model_status = gr.Textbox(label="主模型狀態", value="未載入", interactive=False)
+            model_status = gr.Textbox(label="系統狀態", value="未載入", interactive=False)
             
-            gr.Markdown("### 2. LoRA (選用)")
-            lora_url = gr.Textbox(label="LoRA 下載網址", placeholder="輸入 Civitai 直連下載網址...")
-            lora_scale = gr.Slider(0.1, 2.0, value=0.8, step=0.05, label="LoRA 權重 (Scale)")
+            gr.Markdown("### 2. 自訂 LoRA (選用)")
+            lora_url = gr.Textbox(label="LoRA 下載網址", placeholder="輸入 Civitai 連結...")
+            lora_scale = gr.Slider(0.1, 2.0, value=0.8, step=0.05, label="權重 (Scale)")
             with gr.Row():
-                add_lora_btn = gr.Button("➕ 加入 LoRA")
-                clear_lora_btn = gr.Button("🗑️ 清空 LoRA")
-            lora_status = gr.Textbox(label="已啟用 LoRA 列表", value="無", lines=3, interactive=False)
+                add_lora_btn = gr.Button("➕ 加入")
+                clear_lora_btn = gr.Button("🗑️ 清空")
+            lora_status = gr.Textbox(label="已啟用清單", value="無", lines=2, interactive=False)
 
+        # ── 右側：生成與設定 ──
         with gr.Column(scale=2):
-            gr.Markdown("### 3. 生成設定")
-            prompt = gr.Textbox(label="Prompt", value="1girl, masterpiece, best quality", lines=3)
-            neg = gr.Textbox(label="Negative Prompt", value="low quality, bad anatomy, worst quality", lines=2)
+            use_lcm = gr.Checkbox(label="⚡ 啟用 LCM 極速模式 (強烈建議開啟)", value=True)
+            gr.Markdown("> *開啟 LCM 時：Steps 建議設 4~6，CFG 建議設 1.0~2.0。*\n> *關閉 LCM 時：Steps 建議設 15~20，CFG 建議設 6.0~7.0。*")
+            
+            prompt = gr.Textbox(label="Prompt", value="a beautiful landscape painting, golden hour, highly detailed, masterpiece", lines=3)
+            neg = gr.Textbox(label="Negative Prompt (LCM 模式下將自動忽略)", value="low quality, bad anatomy, worst quality", lines=1)
+            
             with gr.Row():
-                steps = gr.Slider(1, 25, value=10, step=1, label="Steps (CPU建議低於12)")
-                cfg = gr.Slider(1.0, 15.0, value=7.0, step=0.5, label="CFG")
+                steps = gr.Slider(1, 30, value=5, step=1, label="Steps")
+                cfg = gr.Slider(1.0, 10.0, value=1.5, step=0.5, label="CFG Scale")
                 seed = gr.Number(-1, label="Seed (-1=隨機)", precision=0)
             with gr.Row():
-                width = gr.Dropdown([384, 512, 768, 1024], value=512, label="Width")
-                height = gr.Dropdown([384, 512, 768, 1024], value=512, label="Height")
+                # CPU 建議最高不要超過 512
+                width = gr.Dropdown([384, 448, 512, 768], value=384, label="Width")
+                height = gr.Dropdown([384, 448, 512, 768], value=384, label="Height")
                 
-            gen_btn = gr.Button("✨ 生成圖片", variant="primary")
-            
-            with gr.Row():
-                out_img = gr.Image(label="生成結果", type="pil")
-            out_seed = gr.Number(label="Used Seed", precision=0)
+            gen_btn = gr.Button("✨ 生成圖片", variant="primary", size="lg")
+            gen_status = gr.Textbox(label="生成狀態", interactive=False)
+            out_img = gr.Image(label="生成結果", type="pil")
 
     # ── 4. 事件綁定 ──
-    load_preset_btn.click(
-        fn=handle_preset_model, inputs=[preset_dd], outputs=[model_status]
-    )
-    load_civit_btn.click(
-        fn=handle_civitai_model, inputs=[civit_ckpt_url, civit_token], outputs=[model_status]
-    )
-    add_lora_btn.click(
-        fn=add_lora, inputs=[lora_url, lora_scale, civit_token], outputs=[model_status, lora_status]
-    )
-    clear_lora_btn.click(
-        fn=clear_loras, outputs=[model_status]
-    ).then(
-        fn=update_lora_list, outputs=[lora_status]
-    )
+    load_preset_btn.click(fn=handle_preset_model, inputs=[preset_dd], outputs=[model_status])
+    load_civit_btn.click(fn=handle_civitai_model, inputs=[civit_ckpt_url, civit_token], outputs=[model_status])
+    
+    add_lora_btn.click(fn=add_lora, inputs=[lora_url, lora_scale, civit_token], outputs=[model_status, lora_status])
+    clear_lora_btn.click(fn=clear_loras, outputs=[model_status]).then(fn=update_lora_list, outputs=[lora_status])
+    
     gen_btn.click(
-        fn=generate_image, inputs=[prompt, neg, steps, cfg, seed, width, height], outputs=[out_img, out_seed]
+        fn=generate_image, 
+        inputs=[prompt, neg, steps, cfg, seed, width, height, use_lcm], 
+        outputs=[out_img, gen_status]
     )
 
 demo.queue().launch()