| 1 | import time |
| 2 | |
| 3 | from fastapi import APIRouter, File, Form, HTTPException, Request, UploadFile |
| 4 | from fastapi.responses import StreamingResponse |
| 5 | |
| 6 | from api.dependencies import workflow_engine |
| 7 | from api.routers.files import merge_uploaded_file_into_idea |
| 8 | from api.schemas.project import InterventionRequest, ProjectStartRequest |
| 9 | from api.services.project_helpers import ( |
| 10 | make_cancellation, |
| 11 | make_progress_channel, |
| 12 | stream_workflow_task, |
| 13 | ) |
| 14 | from config import settings |
| 15 | |
| 16 | router = APIRouter(tags=["Workflow"]) |
| 17 | |
| 18 | REQUIRED_MODEL_FIELDS = ( |
| 19 | "llm_model", |
| 20 | "vlm_model", |
| 21 | "image_t2i_model", |
| 22 | "image_it2i_model", |
| 23 | ) |
| 24 | |
| 25 | VIDEO_MODE_TO_MODEL_FIELD = { |
| 26 | "first_frame": "video_first_frame_model", |
| 27 | "start_end_frame": "video_start_end_model", |
| 28 | "reference": "video_reference_model", |
| 29 | } |
| 30 | |
| 31 | |
| 32 | def _active_video_model(values: dict) -> str: |
| 33 | mode = values.get("video_generation_mode") or "first_frame" |
| 34 | model_field = VIDEO_MODE_TO_MODEL_FIELD.get(mode, "video_first_frame_model") |
| 35 | default_attr = { |
| 36 | "video_first_frame_model": "VIDEO_FIRST_FRAME_MODEL", |
| 37 | "video_start_end_model": "VIDEO_START_END_MODEL", |
| 38 | "video_reference_model": "VIDEO_REFERENCE_MODEL", |
| 39 | }.get(model_field, "VIDEO_MODEL") |
| 40 | # Legacy session compatibility: fall back to the old single video_model field when mode-specific fields are absent. |
| 41 | return values.get(model_field) or values.get("video_model") or getattr(settings, default_attr, "") or "" |
| 42 | |
| 43 | |
| 44 | def _require_model_fields(values: dict) -> None: |
| 45 | missing = [field for field in REQUIRED_MODEL_FIELDS if not values.get(field)] |
| 46 | if not _active_video_model(values): |
| 47 | missing.append("video_model") |
| 48 | if missing: |
| 49 | raise HTTPException( |
| 50 | status_code=400, |
| 51 | detail=f"Missing required model configuration: {', '.join(missing)}", |
| 52 | ) |
| 53 | |
| 54 | |
| 55 | @router.post("/api/project/start") |
| 56 | async def start_project(req: ProjectStartRequest): |
| 57 | final_idea = merge_uploaded_file_into_idea(req.idea, req.file_path) |
| 58 | _require_model_fields(req.model_dump()) |
| 59 | |
| 60 | session_id = str(int(time.time() * 1000)) |
| 61 | meta = { |
| 62 | "idea": final_idea, |
| 63 | "user_textbox_input": req.idea, |
| 64 | "style": req.style or getattr(settings, "STYLE", None) or "realistic", |
| 65 | "video_ratio": req.video_ratio or "9:16", |
| 66 | "video_resolution": req.video_resolution or "720P", |
| 67 | "expand_idea": req.expand_idea if req.expand_idea is not None else True, |
| 68 | "llm_model": req.llm_model, |
| 69 | "vlm_model": req.vlm_model, |
| 70 | "image_t2i_model": req.image_t2i_model, |
| 71 | "image_it2i_model": req.image_it2i_model, |
| 72 | # Legacy request/session compatibility: clients created before the split only send video_model. |
| 73 | "video_first_frame_model": req.video_first_frame_model or req.video_model or getattr(settings, "VIDEO_FIRST_FRAME_MODEL", ""), |
| 74 | "video_start_end_model": req.video_start_end_model or getattr(settings, "VIDEO_START_END_MODEL", ""), |
| 75 | "video_reference_model": req.video_reference_model or getattr(settings, "VIDEO_REFERENCE_MODEL", ""), |
| 76 | "video_generation_mode": req.video_generation_mode or getattr(settings, "VIDEO_GENERATION_MODE", "first_frame"), |
| 77 | "video_model": req.video_model, |
| 78 | "enable_concurrency": req.enable_concurrency if req.enable_concurrency is not None else True, |
| 79 | "web_search": req.web_search if req.web_search is not None else False, |
| 80 | "episodes": req.episodes if req.episodes is not None else 4, |
| 81 | } |
| 82 | meta["video_model"] = _active_video_model(meta) |
| 83 | session = workflow_engine.create_session(session_id, meta) |
| 84 | |
| 85 | return { |
| 86 | "session_id": session_id, |
| 87 | "status": session["status"], |
| 88 | "params": { |
| 89 | "idea": final_idea, |
| 90 | "file_path": req.file_path, |
| 91 | "style": req.style, |
| 92 | "llm_model": meta["llm_model"], |
| 93 | "vlm_model": meta["vlm_model"], |
| 94 | "image_t2i_model": meta["image_t2i_model"], |
| 95 | "image_it2i_model": meta["image_it2i_model"], |
| 96 | "video_first_frame_model": meta["video_first_frame_model"], |
| 97 | "video_start_end_model": meta["video_start_end_model"], |
| 98 | "video_reference_model": meta["video_reference_model"], |
| 99 | "video_generation_mode": meta["video_generation_mode"], |
| 100 | "video_model": meta["video_model"], |
| 101 | "episodes": meta["episodes"], |
| 102 | "video_ratio": meta["video_ratio"], |
| 103 | "video_resolution": meta["video_resolution"], |
| 104 | } |
| 105 | } |
| 106 | |
| 107 | |
| 108 | @router.post("/api/project/{session_id}/execute/{stage}") |
| 109 | async def execute_stage(session_id: str, stage: str, request: Request): |
| 110 | try: |
| 111 | body = await request.json() |
| 112 | except Exception: |
| 113 | body = {} |
| 114 | |
| 115 | state, input_data = workflow_engine.prepare_stage_execution(session_id, stage, body) |
| 116 | _require_model_fields(input_data) |
| 117 | |
| 118 | cancellation_check, on_disconnect = make_cancellation(workflow_engine, session_id) |
| 119 | progress_events, event_trigger, progress_callback = make_progress_channel() |
| 120 | |
| 121 | return StreamingResponse( |
| 122 | stream_workflow_task( |
| 123 | request=request, |
| 124 | workflow_engine=workflow_engine, |
| 125 | state=state, |
| 126 | stage=stage, |
| 127 | input_data=input_data, |
| 128 | cancellation_check=cancellation_check, |
| 129 | progress_callback=progress_callback, |
| 130 | progress_events=progress_events, |
| 131 | event_trigger=event_trigger, |
| 132 | include_payload_summary=True, |
| 133 | on_disconnect=on_disconnect, |
| 134 | ), |
| 135 | media_type="text/event-stream", |
| 136 | headers={ |
| 137 | "Cache-Control": "no-cache, no-transform", |
| 138 | "X-Accel-Buffering": "no", |
| 139 | "Connection": "keep-alive", |
| 140 | }, |
| 141 | ) |
| 142 | |
| 143 | |
| 144 | @router.get("/api/project/{session_id}/status") |
| 145 | async def get_project_status(session_id: str): |
| 146 | snapshot = workflow_engine.get_status_snapshot(session_id) |
| 147 | if not snapshot: |
| 148 | raise HTTPException(404, "Session not found") |
| 149 | return snapshot |
| 150 | |
| 151 | |
| 152 | @router.get("/api/project/{session_id}/status/from_disk") |
| 153 | async def get_project_status_from_disk(session_id: str): |
| 154 | # 兼容旧前端路由名；实际读取统一走 WorkflowEngine 的内存状态入口。 |
| 155 | snapshot = workflow_engine.get_status_snapshot(session_id) |
| 156 | if not snapshot: |
| 157 | raise HTTPException(404, "Session not found") |
| 158 | return snapshot |
| 159 | |
| 160 | |
| 161 | @router.get("/api/project/{session_id}/artifact/{stage}") |
| 162 | async def get_artifact(session_id: str, stage: str): |
| 163 | try: |
| 164 | artifact = workflow_engine.get_artifact_snapshot(session_id, stage) |
| 165 | except KeyError: |
| 166 | raise HTTPException(404, "Session not found") |
| 167 | |
| 168 | if artifact is not None: |
| 169 | return {"stage": stage, "artifact": artifact} |
| 170 | |
| 171 | raise HTTPException(404, f"Artifact for stage '{stage}' not found") |
| 172 | |
| 173 | |
| 174 | @router.patch("/api/project/{session_id}/models") |
| 175 | async def update_models(session_id: str, request: Request): |
| 176 | body = await request.json() |
| 177 | allowed_keys = ( |
| 178 | "llm_model", |
| 179 | "vlm_model", |
| 180 | "image_t2i_model", |
| 181 | "image_it2i_model", |
| 182 | "video_model", |
| 183 | "video_first_frame_model", |
| 184 | "video_start_end_model", |
| 185 | "video_reference_model", |
| 186 | "video_generation_mode", |
| 187 | "video_ratio", |
| 188 | "video_resolution", |
| 189 | "style", |
| 190 | "enable_concurrency", |
| 191 | ) |
| 192 | try: |
| 193 | return workflow_engine.update_session_meta(session_id, body if isinstance(body, dict) else {}, allowed_keys) |
| 194 | except KeyError: |
| 195 | raise HTTPException(404, "Session not found") |
| 196 | |
| 197 | |
| 198 | @router.post("/api/project/{session_id}/artifact/{stage}/upload_image") |
| 199 | async def upload_artifact_image( |
| 200 | session_id: str, |
| 201 | stage: str, |
| 202 | item_type: str = Form(...), |
| 203 | item_id: str = Form(...), |
| 204 | file: UploadFile = File(...), |
| 205 | ): |
| 206 | """Upload a user-provided image into a stage artifact and persist the session.""" |
| 207 | try: |
| 208 | return workflow_engine.upload_artifact_image( |
| 209 | session_id=session_id, |
| 210 | stage=stage, |
| 211 | item_type=item_type, |
| 212 | item_id=item_id, |
| 213 | file_obj=file.file, |
| 214 | filename=file.filename or "", |
| 215 | ) |
| 216 | except KeyError: |
| 217 | raise HTTPException(404, "Session not found") |
| 218 | except ValueError as exc: |
| 219 | raise HTTPException(400, detail=str(exc)) from exc |
| 220 | except RuntimeError as exc: |
| 221 | raise HTTPException(500, detail=str(exc)) from exc |
| 222 | |
| 223 | |
| 224 | @router.patch("/api/project/{session_id}/artifact/{stage}") |
| 225 | async def update_artifact(session_id: str, stage: str, request: Request): |
| 226 | """保存用户在某阶段的选择/修改，同时更新内存状态和磁盘快照。""" |
| 227 | body = await request.json() |
| 228 | try: |
| 229 | return workflow_engine.update_artifact(session_id, stage, body if isinstance(body, dict) else {}) |
| 230 | except KeyError: |
| 231 | raise HTTPException(404, "Session not found") |
| 232 | |
| 233 | |
| 234 | @router.post("/api/project/{session_id}/intervene") |
| 235 | async def intervene(session_id: str, req: InterventionRequest, request: Request): |
| 236 | try: |
| 237 | state, input_data = workflow_engine.prepare_intervention_execution( |
| 238 | session_id=session_id, |
| 239 | stage=req.stage, |
| 240 | modifications=req.modifications, |
| 241 | ) |
| 242 | except KeyError: |
| 243 | raise HTTPException(404, "Session not found") |
| 244 | |
| 245 | cancellation_check, on_disconnect = make_cancellation(workflow_engine, session_id) |
| 246 | progress_events, event_trigger, progress_callback = make_progress_channel() |
| 247 | |
| 248 | return StreamingResponse( |
| 249 | stream_workflow_task( |
| 250 | request=request, |
| 251 | workflow_engine=workflow_engine, |
| 252 | state=state, |
| 253 | stage=req.stage, |
| 254 | input_data=input_data, |
| 255 | cancellation_check=cancellation_check, |
| 256 | progress_callback=progress_callback, |
| 257 | progress_events=progress_events, |
| 258 | event_trigger=event_trigger, |
| 259 | intervention=req.modifications, |
| 260 | on_disconnect=on_disconnect, |
| 261 | ), |
| 262 | media_type="text/event-stream", |
| 263 | headers={ |
| 264 | "Cache-Control": "no-cache, no-transform", |
| 265 | "X-Accel-Buffering": "no", |
| 266 | "Connection": "keep-alive", |
| 267 | }, |
| 268 | ) |
| 269 | |
| 270 | |
| 271 | @router.post("/api/project/{session_id}/continue") |
| 272 | async def continue_workflow(session_id: str): |
| 273 | if not workflow_engine.get_status_snapshot(session_id): |
| 274 | raise HTTPException(404, "Session not found") |
| 275 | return await workflow_engine.continue_workflow(session_id) |
| 276 | |
| 277 | |
| 278 | @router.post("/api/project/{session_id}/stop") |
| 279 | async def stop_project(session_id: str): |
| 280 | workflow_engine.stop_session(session_id) |
| 281 | return {"status": "stopped", "session_id": session_id} |
| 282 | |
| 283 | |
| 284 | @router.get("/api/project/{session_id}/scene/{scene_number}/assets") |
| 285 | async def check_scene_assets(session_id: str, scene_number: int): |
| 286 | try: |
| 287 | return workflow_engine.get_scene_asset_counts(session_id, scene_number) |
| 288 | except KeyError: |
| 289 | raise HTTPException(404, "Session not found") |
| 290 |