Update app.py

app.py

@@ -648,16 +648,27 @@ class GenerateRequest(BaseModel):
     char_strength: float = 1.0
     base_lora_states: list = []
 
+import threading
+# 专门用于保护 tasks_status 字典的状态锁
+_status_lock = threading.RLock()
 _model_switch_lock = threading.Lock()
 
 def background_generate_task(task_id: str, req: GenerateRequest, ip_token: str = None, request: Request = None):
     global _last_used_architecture, _last_used_unet
 
     _current_generation_task_id.set(task_id)
-    tasks_status[task_id] = {"status": "processing", "message": "⏳ Queue for GPU..."}
+    
+    # 初始状态
+    with _status_lock:
+        tasks_status[task_id] = {"status": "processing", "message": "⏳ Queue for GPU..."}
+        
     print(f"[Token探针] task={task_id[:8]} | route_token_present={bool(ip_token)}", flush=True)
 
     try:
+        # --- 阶段 1：环境准备与模型下载 ---
+        with _status_lock:
+            tasks_status[task_id]["message"] = "📥 Downloading LoRA / Preparing environment..."
+            
         download_start = time.time()
         prepare_env(req.char_enable, req.char_lora, req.unet_choice, task_id=task_id)
         real_download_time = round(time.time() - download_start, 2)
@@ -668,20 +679,24 @@ def background_generate_task(task_id: str, req: GenerateRequest, ip_token: str =
         if need_heavy_load:
             with _model_switch_lock:
                 if _last_used_architecture != current_arch:
-                    tasks_status[task_id]["message"] = f"📦 Queue for GPU... (Loading Base Model: {current_arch.upper()})"
+                    with _status_lock:
+                        tasks_status[task_id]["message"] = f"📦 Loading Base Model: {current_arch.upper()}..."
                     _last_used_architecture = current_arch
                     _last_used_unet = req.unet_choice
                 elif _last_used_unet != req.unet_choice:
-                    tasks_status[task_id]["message"] = "🔄 Queue for GPU... (Swapping Model Checkpoint)"
+                    with _status_lock:
+                        tasks_status[task_id]["message"] = "🔄 Swapping Model Checkpoint..."
                     _last_used_unet = req.unet_choice
         else:
-            tasks_status[task_id]["message"] = "🎨 Rendering image (GPU processing)..."
+            with _status_lock:
+                tasks_status[task_id]["message"] = "🎨 GPU processing..."
 
         final_seed = req.seed if req.seed >= 0 else random.randint(0, 2**32 - 1)
-        tasks_status[task_id]["seed"] = final_seed
+        with _status_lock:
+            tasks_status[task_id]["seed"] = final_seed
+            tasks_status[task_id]["queue_start_time"] = time.time()
+            tasks_status[task_id]["message"] = "⏳ Requesting GPU resources..."
 
-        tasks_status[task_id]["queue_start_time"] = time.time()
-        tasks_status[task_id]["message"] = "⏳ Requesting GPU resources..."
         print(f"[GPU 探针] 🔴 发起 Request，申请 GPU 资源...", flush=True)
 
         ctx = contextvars.copy_context()
@@ -689,15 +704,6 @@ def background_generate_task(task_id: str, req: GenerateRequest, ip_token: str =
         def _run_generate():
             if request is not None:
                 LocalContext.request.set(request)
-
-            _req_in_ctx = LocalContext.request.get(None)
-            _token_in_ctx = _req_in_ctx.headers.get("x-ip-token") if _req_in_ctx is not None else None
-            print(
-                f"[Token探针] task={task_id[:8]} | ctx_request_present={_req_in_ctx is not None}"
-                f" | ctx_token_present={bool(_token_in_ctx)}",
-                flush=True,
-            )
-
             return generate(
                 req.prompt, req.unet_choice, req.custom_w, req.custom_h,
                 req.steps, req.cfg, final_seed, req.sampler, req.scheduler,
@@ -706,27 +712,18 @@ def background_generate_task(task_id: str, req: GenerateRequest, ip_token: str =
                 task_id=task_id,
             )
 
+        # --- 阶段 2：核心执行（阻塞式，前端此时只看到进度条走动） ---
         image_path, real_workflow = ctx.run(_run_generate)
 
-        # 🌟 重点：把真正的 workflow 转成字符串
-        real_workflow_str = json.dumps(real_workflow, ensure_ascii=False)
-
-        tasks_status[task_id]["message"] = "✨ Processing output..."
-        
+        # --- 阶段 3：死磕后处理（前端此时处于静默等待，避免出现“快好了”的错觉） ---
         from PIL import Image
         
-        # 提取文件名，并确保存成 webp
         original_filename = os.path.basename(image_path)
         webp_filename = original_filename.replace(".png", ".webp")
-        
-        # 🌟 核心修改 1：定义物理保存路径
         output_webp_path = os.path.join(OUTPUT, webp_filename)
 
         with Image.open(image_path) as img:
-            # 1. 获取原始元数据（ComfyUI 存的 PNG 里本身就带完整的 JSON 数据）
             metadata = img.info
-            
-            # 2. 构建标准的 EXIF 标签（���全照抄 GitHub 完美方案）
             PROMPT_TAG = 0x0110
             EXTRA_METADATA_TAG = 0x010f
             output_metadata = img.getexif()
@@ -738,11 +735,10 @@ def background_generate_task(task_id: str, req: GenerateRequest, ip_token: str =
                     output_metadata[EXTRA_METADATA_TAG] = "Workflow:" + value
                     EXTRA_METADATA_TAG -= 1
 
-            # 3. 转换 RGB（防止 WebP 保存报错）
             if img.mode in ("RGBA", "P"):
                 img = img.convert("RGB")
 
-            # 4. 🌟 核心修改 2：直接保存到磁盘，不再写入内存
+            # 🌟 性能妥协：method=4, quality=90，速度极快，肉眼无感
             img.save(
                 output_webp_path, 
                 "WEBP", 
@@ -752,16 +748,14 @@ def background_generate_task(task_id: str, req: GenerateRequest, ip_token: str =
                 exif=output_metadata
             )
 
-        # 删除原来的 PNG，节省空间
         try:
             os.remove(image_path)
         except Exception:
             pass
 
-        # 🌟 核心修改 3：构建前端可以直接访问的 URL
-        # 注意：前端/后端如果在同一个根目录，用 /output/ 即可。如果跨域，请换成你的 API_BASE_URL
         image_url = f"/output/{webp_filename}"
 
+        # --- 阶段 4：提取指标 ---
         final_queue = 0
         final_workflow = 0
         final_load_time = 0
@@ -775,26 +769,30 @@ def background_generate_task(task_id: str, req: GenerateRequest, ip_token: str =
                     final_workflow = saved_metrics.get("workflow", 0)
                     final_load_time = saved_metrics.get("load_time", 0)
                 os.remove(metric_file)
-            except Exception as e:
-                print(f"读取指标文件失败: {e}")
-
-        tasks_status[task_id] = {
-            "status": "success",
-            "image": image_url,  # 👈 核心修改 4：直接把 URL 传回去
-            "filename": webp_filename,
-            "seed": final_seed,
-            "metrics": {
-                "download": real_download_time,
-                "queue": final_queue,
-                "workflow": final_workflow,
-                "load_time": final_load_time,
-                "steps": req.steps,
-                "resolved_seed": final_seed 
-            },
-            "anonymous": ip_token is None,
-        }
+            except Exception:
+                pass
+
+        # --- 阶段 5：高光时刻！直接抛出最终结果，抛弃一切中间态日志 ---
+        with _status_lock:
+            tasks_status[task_id] = {
+                "status": "success",
+                "image": image_url, 
+                "filename": webp_filename,
+                "seed": final_seed,
+                "metrics": {
+                    "download": real_download_time,
+                    "queue": final_queue,
+                    "workflow": final_workflow,
+                    "load_time": final_load_time,
+                    "steps": req.steps,
+                    "resolved_seed": final_seed 
+                },
+                "anonymous": ip_token is None,
+            }
+            
     except Exception as e:
-        tasks_status[task_id] = {"status": "error", "message": str(e)}
+        with _status_lock:
+            tasks_status[task_id] = {"status": "error", "message": str(e)}
     finally:
         _current_generation_task_id.set(None)