Update app.py

app.py

@@ -10,7 +10,7 @@ import shutil
 import asyncio
 import requests
 import httpx
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Optional
 from threading import Lock
 from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, UploadFile, File
 from fastapi.staticfiles import StaticFiles
@@ -32,18 +32,31 @@ app.add_middleware(
 class ConnectionManager:
     def __init__(self):
         self.active_connections: List[WebSocket] = []
+        self.user_connections: Dict[str, WebSocket] = {}
 
-    async def connect(self, websocket: WebSocket):
+    async def connect(self, websocket: WebSocket, client_id: str = None):
         await websocket.accept()
         self.active_connections.append(websocket)
+        if client_id:
+            self.user_connections[client_id] = websocket
         print(f"WS Connected. Total: {len(self.active_connections)}")
         await self.broadcast_count()
 
-    async def disconnect(self, websocket: WebSocket):
+    async def disconnect(self, websocket: WebSocket, client_id: str = None):
         if websocket in self.active_connections:
             self.active_connections.remove(websocket)
-            print(f"WS Disconnected. Total: {len(self.active_connections)}")
-            await self.broadcast_count()
+        if client_id and client_id in self.user_connections:
+            del self.user_connections[client_id]
+        print(f"WS Disconnected. Total: {len(self.active_connections)}")
+        await self.broadcast_count()
+
+    async def send_personal_message(self, message: dict, client_id: str):
+        if client_id in self.user_connections:
+            try:
+                await self.user_connections[client_id].send_text(json.dumps(message))
+            except Exception as e:
+                print(f"WS Send Error ({client_id}): {e}")
+                self.disconnect(self.user_connections[client_id], client_id)
 
     async def broadcast_count(self):
         count = len(self.active_connections)
@@ -79,8 +92,8 @@ async def startup_event():
     GLOBAL_LOOP = asyncio.get_running_loop()
 
 @app.websocket("/ws/stats")
-async def websocket_endpoint(websocket: WebSocket):
-    await manager.connect(websocket)
+async def websocket_endpoint(websocket: WebSocket, client_id: str = None):
+    await manager.connect(websocket, client_id)
     try:
         while True:
             # 接收客户端心跳包
@@ -89,10 +102,10 @@ async def websocket_endpoint(websocket: WebSocket):
                 await websocket.send_text(json.dumps({"type": "pong"}))
     except WebSocketDisconnect:
         print(f"WebSocket disconnected normally: {id(websocket)}")
-        await manager.disconnect(websocket)
+        await manager.disconnect(websocket, client_id)
     except Exception as e:
         print(f"WS Error for {id(websocket)}: {e}")
-        await manager.disconnect(websocket)
+        await manager.disconnect(websocket, client_id)
 
 # --- 配置区域 ---
 # 支持多卡负载均衡：配置多个 ComfyUI 地址
@@ -144,10 +157,12 @@ class GenerateRequest(BaseModel):
 
 class CloudGenRequest(BaseModel):
     prompt: str
-    api_key: str = ""
+    api_key: str
     resolution: str = "1024x1024"
-    client_id: str = "default"
-    type: str = "zimage"
+    client_id: Optional[str] = None
+    type: str = "default"
+    image_urls: List[str] = []
+    model: str = ""
 
 class DeleteHistoryRequest(BaseModel):
     timestamp: float
@@ -320,6 +335,14 @@ async def get_history_api(type: str = None):
                     return 0 # 旧数据排在最后
                 
                 data.sort(key=sort_key, reverse=True)
+
+                # 补充 is_cloud 字段：如果历史记录中没有标记，但文件名包含特征字符，则补充标记
+                for item in data:
+                    if "is_cloud" not in item and item.get("images"):
+                        # 检查是否有任意一张图片符合 cloud 特征
+                        if any("cloud_angle" in img or "cloud_" in img for img in item["images"]):
+                            item["is_cloud"] = True
+
                 return data
         except Exception as e:
             print(f"读取历史文件失败: {e}")
@@ -436,6 +459,267 @@ async def delete_global_token():
             except: pass
     return {"success": True}
 
+class CloudPollRequest(BaseModel):
+    task_id: str
+    api_key: str
+    client_id: Optional[str] = None
+
+@app.post("/api/angle/poll_status")
+async def poll_angle_cloud(req: CloudPollRequest):
+    """
+    Resume polling for an existing Angle task.
+    """
+    base_url = 'https://api-inference.modelscope.cn/'
+    clean_token = req.api_key.strip()
+    
+    headers = {
+        "Authorization": f"Bearer {clean_token}",
+        "Content-Type": "application/json",
+        "X-ModelScope-Async-Mode": "true"
+    }
+    
+    task_id = req.task_id
+    print(f"Resuming polling for Angle Task: {task_id}")
+
+    try:
+        async with httpx.AsyncClient(timeout=30) as client:
+            # Poll Status (Another 300 retries)
+            for i in range(300):
+                await asyncio.sleep(2) 
+                try:
+                    result = await client.get(
+                        f"{base_url}v1/tasks/{task_id}",
+                        headers={**headers, "X-ModelScope-Task-Type": "image_generation"},
+                    )
+                    data = result.json()
+                    status = data.get("task_status")
+
+                    if status == "SUCCEED":
+                        img_url = data["output_images"][0]
+                        print(f"Angle Task SUCCEED: {img_url}")
+                        
+                        if req.client_id:
+                             await manager.send_personal_message({
+                                 "type": "cloud_status", 
+                                 "status": "SUCCEED", 
+                                 "task_id": task_id
+                             }, req.client_id)
+
+                        # Download logic
+                        local_path = ""
+                        try:
+                            async with httpx.AsyncClient() as dl_client:
+                                img_res = await dl_client.get(img_url)
+                                if img_res.status_code == 200:
+                                    filename = f"cloud_angle_{int(time.time())}.png"
+                                    file_path = os.path.join(OUTPUT_DIR, filename)
+                                    with open(file_path, "wb") as f:
+                                        f.write(img_res.content)
+                                    local_path = f"/output/{filename}"
+                                else:
+                                    local_path = img_url
+                        except Exception:
+                            local_path = img_url
+
+                        record = {
+                            "timestamp": time.time(),
+                            "prompt": f"Resumed {task_id}",
+                            "images": [local_path],
+                            "type": "angle"
+                        }
+                        save_to_history(record)
+                        return {"url": local_path}
+
+                    elif status == "FAILED":
+                        if req.client_id:
+                            await manager.send_personal_message({
+                                "type": "cloud_status", 
+                                "status": "FAILED", 
+                                "task_id": task_id
+                            }, req.client_id)
+                        raise Exception(f"ModelScope task failed: {data}")
+                    
+                    if i % 5 == 0:
+                        print(f"Angle Task {task_id} status: {status} ({i}/150)")
+                        if req.client_id:
+                            await manager.send_personal_message({
+                                "type": "cloud_status", 
+                                "status": f"{status} ({i}/150)", 
+                                "task_id": task_id,
+                                "progress": i,
+                                "total": 150
+                            }, req.client_id)
+                        
+                except Exception as loop_e:
+                    print(f"Angle polling error: {loop_e}")
+                    continue
+
+            print(f"Angle Task Timeout Again: {task_id}")
+            if req.client_id:
+                await manager.send_personal_message({
+                    "type": "cloud_status", 
+                    "status": "TIMEOUT", 
+                    "task_id": task_id
+                }, req.client_id)
+
+            return {"status": "timeout", "task_id": task_id, "message": "Task still pending"}
+
+    except Exception as e:
+        print(f"Angle polling error: {e}")
+        raise HTTPException(status_code=400, detail=str(e))
+
+@app.post("/api/angle/generate")
+async def generate_angle_cloud(req: CloudGenRequest):
+    """
+    Dedicated endpoint for Angle/Qwen-Image-Edit tasks.
+    Logic mirrors test/main.py but uses async httpx.
+    """
+    base_url = 'https://api-inference.modelscope.cn/'
+    clean_token = req.api_key.strip()
+    
+    headers = {
+        "Authorization": f"Bearer {clean_token}",
+        "Content-Type": "application/json",
+        "X-ModelScope-Async-Mode": "true"
+    }
+
+    # Prepare payload exactly as in test/main.py
+    # test/main.py: "image_url": [data_uri]
+    # req.image_urls is already a list of strings
+    payload = {
+        "model": "Qwen/Qwen-Image-Edit-2511",
+        "prompt": req.prompt.strip(),
+        "image_url": req.image_urls 
+    }
+
+    print(f"Angle Cloud Request: {payload['model']}, Prompt: {payload['prompt'][:20]}...")
+
+    try:
+        async with httpx.AsyncClient(timeout=30) as client:
+            # 1. Submit Task
+            submit_res = await client.post(
+                f"{base_url}v1/images/generations",
+                headers=headers,
+                json=payload # httpx handles json serialization
+            )
+
+            if submit_res.status_code != 200:
+                try:
+                    detail = submit_res.json()
+                except:
+                    detail = submit_res.text
+                print(f"Angle Submit Error: {detail}")
+                raise HTTPException(status_code=submit_res.status_code, detail=detail)
+
+            task_id = submit_res.json().get("task_id")
+            print(f"Angle Task Submitted, ID: {task_id}")
+
+            # Notify frontend via WS
+            if req.client_id:
+                await manager.send_personal_message({
+                    "type": "cloud_status", 
+                    "status": "SUBMITTED", 
+                    "task_id": task_id,
+                    "progress": 0,
+                    "total": 150
+                }, req.client_id)
+
+            # 2. Poll Status (300 retries * 2s = 600s / 10min)
+            for i in range(300):
+                await asyncio.sleep(2) 
+                try:
+                    result = await client.get(
+                        f"{base_url}v1/tasks/{task_id}",
+                        headers={**headers, "X-ModelScope-Task-Type": "image_generation"},
+                    )
+                    data = result.json()
+                    status = data.get("task_status")
+
+                    if status == "SUCCEED":
+                        img_url = data["output_images"][0]
+                        print(f"Angle Task SUCCEED: {img_url}")
+                        
+                        # Notify WS success
+                        if req.client_id:
+                            await manager.send_personal_message({
+                                "type": "cloud_status", 
+                                "status": "SUCCEED", 
+                                "task_id": task_id
+                            }, req.client_id)
+
+                        # Download and Save Logic (reused from original generate)
+                        local_path = ""
+                        try:
+                             # 异步下载
+                            async with httpx.AsyncClient() as dl_client:
+                                img_res = await dl_client.get(img_url)
+                                if img_res.status_code == 200:
+                                    filename = f"cloud_angle_{int(time.time())}.png"
+                                    file_path = os.path.join(OUTPUT_DIR, filename)
+                                    with open(file_path, "wb") as f:
+                                        f.write(img_res.content)
+                                    local_path = f"/output/{filename}"
+                                    print(f"Angle Image saved: {local_path}")
+                                else:
+                                    local_path = img_url
+                        except Exception as dl_e:
+                            print(f"Download error: {dl_e}")
+                            local_path = img_url
+
+                        # Save history
+                        record = {
+                            "timestamp": time.time(),
+                            "prompt": req.prompt,
+                            "images": [local_path],
+                            "type": "angle", # Distinct type
+                            "is_cloud": True
+                        }
+                        save_to_history(record)
+                        return {"url": local_path}
+
+                    elif status == "FAILED":
+                        if req.client_id:
+                            await manager.send_personal_message({
+                                "type": "cloud_status", 
+                                "status": "FAILED", 
+                                "task_id": task_id
+                            }, req.client_id)
+                        raise Exception(f"ModelScope task failed: {data}")
+                    
+                    # Log polling status every 5 times (10 seconds)
+                    if i % 5 == 0:
+                        print(f"Angle Task {task_id} status: {status} ({i}/150)")
+                        if req.client_id:
+                            await manager.send_personal_message({
+                                "type": "cloud_status", 
+                                "status": f"{status} ({i}/150)", 
+                                "task_id": task_id,
+                                "progress": i,
+                                "total": 150
+                            }, req.client_id)
+                        
+                except Exception as loop_e:
+                    # Log polling errors
+                    print(f"Angle polling error (retrying): {loop_e}")
+                    continue
+
+            # Timeout Handling
+            print(f"Angle Task Timeout: {task_id}")
+            if req.client_id:
+                await manager.send_personal_message({
+                    "type": "cloud_status", 
+                    "status": "TIMEOUT", 
+                    "task_id": task_id
+                }, req.client_id)
+                
+            # Instead of raising Exception, return special status
+            return {"status": "timeout", "task_id": task_id, "message": "Task still pending"}
+
+    except Exception as e:
+        print(f"Angle generation error: {e}")
+        raise HTTPException(status_code=400, detail=str(e))
+
+
 @app.post("/generate")
 async def generate_cloud(req: CloudGenRequest):
     base_url = 'https://api-inference.modelscope.cn/'
@@ -447,6 +731,11 @@ async def generate_cloud(req: CloudGenRequest):
     }
 
     # 按照官方 Z-Image 标准版参数
+    # if req.type == "angle" or req.model == "Qwen/Qwen-Image-Edit-2511":
+    #     # Deprecated: Angle logic moved to /api/angle/generate
+    #     pass
+    
+    # 默认 Z-Image Turbo 模式
     payload = {
         "model": "Tongyi-MAI/Z-Image-Turbo", 
         "prompt": req.prompt.strip(),
@@ -458,10 +747,13 @@ async def generate_cloud(req: CloudGenRequest):
         async with httpx.AsyncClient(timeout=30) as client:
             # A. 提交异步任务
             print(f"Submitting ModelScope task for prompt: {req.prompt[:20]}...")
+            
+            # Use json parameter to ensure standard serialization (matching requests behavior)
+            # This handles Content-Type and ensure_ascii default behavior
             submit_res = await client.post(
                 f"{base_url}v1/images/generations",
                 headers={**headers, "X-ModelScope-Async-Mode": "true"},
-                content=json.dumps(payload, ensure_ascii=False).encode('utf-8')
+                json=payload
             )
             
             if submit_res.status_code != 200:
@@ -477,8 +769,8 @@ async def generate_cloud(req: CloudGenRequest):
             print(f"Task submitted, ID: {task_id}")
 
             # B. 轮询任务状态
-            # 增加到 60 次轮询 * 3秒 = 180秒 (3分钟) 超时
-            for i in range(60):
+            # 增加到 200 次轮询 * 3秒 = 600秒 (10分钟) 超时
+            for i in range(200):
                 await asyncio.sleep(3)
                 try:
                     result = await client.get(