返回 Pixelle-Video
persistence.py
根目录 / pixelle_video / services / persistence.py
1 # Copyright (C) 2025 AIDC-AI
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 # http://www.apache.org/licenses/LICENSE-2.0
7 # Unless required by applicable law or agreed to in writing, software
8 # distributed under the License is distributed on an "AS IS" BASIS,
9 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10 # See the License for the specific language governing permissions and
11 # limitations under the License.
12
13 """
14 Persistence Service
15
16 Handles task metadata and storyboard persistence to filesystem.
17 """
18
19 import json
20 from pathlib import Path
21 from typing import List, Optional, Dict, Any
22 from datetime import datetime
23 from loguru import logger
24
25 from pixelle_video.models.storyboard import Storyboard, StoryboardFrame, StoryboardConfig, ContentMetadata
26
27
28 class PersistenceService:
29 """
30 Task persistence service using filesystem (JSON)
31
32 File structure:
33 output/
34 └── {task_id}/
35 ├── metadata.json # Task metadata (input, result, config)
36 ├── storyboard.json # Storyboard data (frames, prompts)
37 ├── final.mp4
38 └── frames/
39 ├── 01_audio.mp3
40 ├── 01_image.png
41 └── ...
42
43 Usage:
44 persistence = PersistenceService()
45
46 # Save metadata
47 await persistence.save_task_metadata(task_id, metadata)
48
49 # Save storyboard
50 await persistence.save_storyboard(task_id, storyboard)
51
52 # Load task
53 metadata = await persistence.load_task_metadata(task_id)
54 storyboard = await persistence.load_storyboard(task_id)
55
56 # List all tasks
57 tasks = await persistence.list_tasks(status="completed", limit=50)
58 """
59
60 def __init__(self, output_dir: str = "output"):
61 """
62 Initialize persistence service
63
64 Args:
65 output_dir: Base output directory (default: "output")
66 """
67 self.output_dir = Path(output_dir)
68 self.output_dir.mkdir(exist_ok=True)
69
70 # Index file for fast listing
71 self.index_file = self.output_dir / ".index.json"
72 self._ensure_index()
73
74 def get_task_dir(self, task_id: str) -> Path:
75 """Get task directory path"""
76 return self.output_dir / task_id
77
78 def get_metadata_path(self, task_id: str) -> Path:
79 """Get metadata.json path"""
80 return self.get_task_dir(task_id) / "metadata.json"
81
82 def get_storyboard_path(self, task_id: str) -> Path:
83 """Get storyboard.json path"""
84 return self.get_task_dir(task_id) / "storyboard.json"
85
86 # ========================================================================
87 # Metadata Operations
88 # ========================================================================
89
90 async def save_task_metadata(
91 self,
92 task_id: str,
93 metadata: Dict[str, Any]
94 ):
95 """
96 Save task metadata to filesystem
97
98 Args:
99 task_id: Task ID
100 metadata: Metadata dict with structure:
101 {
102 "task_id": str,
103 "created_at": str,
104 "completed_at": str (optional),
105 "status": str,
106 "input": dict,
107 "result": dict (optional),
108 "config": dict
109 }
110 """
111 try:
112 task_dir = self.get_task_dir(task_id)
113 task_dir.mkdir(parents=True, exist_ok=True)
114
115 metadata_path = self.get_metadata_path(task_id)
116
117 # Ensure task_id is set
118 metadata["task_id"] = task_id
119
120 # Convert datetime objects to ISO format strings
121 if "created_at" in metadata and isinstance(metadata["created_at"], datetime):
122 metadata["created_at"] = metadata["created_at"].isoformat()
123 if "completed_at" in metadata and isinstance(metadata["completed_at"], datetime):
124 metadata["completed_at"] = metadata["completed_at"].isoformat()
125
126 with open(metadata_path, "w", encoding="utf-8") as f:
127 json.dump(metadata, f, indent=2, ensure_ascii=False)
128
129 logger.debug(f"Saved task metadata: {task_id}")
130
131 # Update index
132 await self._update_index_for_task(task_id, metadata)
133
134 except Exception as e:
135 logger.error(f"Failed to save task metadata {task_id}: {e}")
136 raise
137
138 async def load_task_metadata(self, task_id: str) -> Optional[Dict[str, Any]]:
139 """
140 Load task metadata from filesystem
141
142 Args:
143 task_id: Task ID
144
145 Returns:
146 Metadata dict or None if not found
147 """
148 try:
149 metadata_path = self.get_metadata_path(task_id)
150
151 if not metadata_path.exists():
152 return None
153
154 with open(metadata_path, "r", encoding="utf-8") as f:
155 metadata = json.load(f)
156
157 return metadata
158
159 except Exception as e:
160 logger.error(f"Failed to load task metadata {task_id}: {e}")
161 return None
162
163 async def update_task_status(
164 self,
165 task_id: str,
166 status: str,
167 error: Optional[str] = None
168 ):
169 """
170 Update task status in metadata
171
172 Args:
173 task_id: Task ID
174 status: New status (pending, running, completed, failed, cancelled)
175 error: Error message (optional, for failed status)
176 """
177 try:
178 metadata = await self.load_task_metadata(task_id)
179 if not metadata:
180 logger.warning(f"Cannot update status: task {task_id} not found")
181 return
182
183 metadata["status"] = status
184
185 if status in ["completed", "failed", "cancelled"]:
186 metadata["completed_at"] = datetime.now().isoformat()
187
188 if error:
189 metadata["error"] = error
190
191 await self.save_task_metadata(task_id, metadata)
192
193 except Exception as e:
194 logger.error(f"Failed to update task status {task_id}: {e}")
195
196 # ========================================================================
197 # Storyboard Operations
198 # ========================================================================
199
200 async def save_storyboard(
201 self,
202 task_id: str,
203 storyboard: Storyboard
204 ):
205 """
206 Save storyboard to filesystem
207
208 Args:
209 task_id: Task ID
210 storyboard: Storyboard instance
211 """
212 try:
213 task_dir = self.get_task_dir(task_id)
214 task_dir.mkdir(parents=True, exist_ok=True)
215
216 storyboard_path = self.get_storyboard_path(task_id)
217
218 # Convert storyboard to dict
219 storyboard_dict = self._storyboard_to_dict(storyboard)
220
221 with open(storyboard_path, "w", encoding="utf-8") as f:
222 json.dump(storyboard_dict, f, indent=2, ensure_ascii=False)
223
224 logger.debug(f"Saved storyboard: {task_id}")
225
226 except Exception as e:
227 logger.error(f"Failed to save storyboard {task_id}: {e}")
228 raise
229
230 async def load_storyboard(self, task_id: str) -> Optional[Storyboard]:
231 """
232 Load storyboard from filesystem
233
234 Args:
235 task_id: Task ID
236
237 Returns:
238 Storyboard instance or None if not found
239 """
240 try:
241 storyboard_path = self.get_storyboard_path(task_id)
242
243 if not storyboard_path.exists():
244 return None
245
246 with open(storyboard_path, "r", encoding="utf-8") as f:
247 storyboard_dict = json.load(f)
248
249 # Convert dict to storyboard
250 storyboard = self._dict_to_storyboard(storyboard_dict)
251
252 return storyboard
253
254 except Exception as e:
255 logger.error(f"Failed to load storyboard {task_id}: {e}")
256 return None
257
258 # ========================================================================
259 # Task Listing & Querying
260 # ========================================================================
261
262 async def list_tasks(
263 self,
264 status: Optional[str] = None,
265 limit: int = 50,
266 offset: int = 0
267 ) -> List[Dict[str, Any]]:
268 """
269 List tasks with optional filtering
270
271 Args:
272 status: Filter by status (pending, running, completed, failed, cancelled)
273 limit: Maximum number of tasks to return
274 offset: Number of tasks to skip
275
276 Returns:
277 List of metadata dicts, sorted by created_at descending
278 """
279 try:
280 index = self._load_index()
281 tasks = index.get("tasks", [])
282
283 # Filter by status
284 if status:
285 tasks = [t for t in tasks if t.get("status") == status]
286
287 # Sort by created_at descending
288 tasks.sort(key=lambda t: t.get("created_at", ""), reverse=True)
289
290 # Apply pagination
291 return tasks[offset:offset + limit]
292
293 except Exception as e:
294 logger.error(f"Failed to list tasks: {e}")
295 return []
296
297 async def task_exists(self, task_id: str) -> bool:
298 """Check if task exists"""
299 return self.get_task_dir(task_id).exists()
300
301 # ========================================================================
302 # Serialization Helpers
303 # ========================================================================
304
305 def _storyboard_to_dict(self, storyboard: Storyboard) -> Dict[str, Any]:
306 """Convert Storyboard to dict for JSON serialization"""
307 return {
308 "title": storyboard.title,
309 "config": self._config_to_dict(storyboard.config),
310 "frames": [self._frame_to_dict(frame) for frame in storyboard.frames],
311 "content_metadata": self._content_metadata_to_dict(storyboard.content_metadata) if storyboard.content_metadata else None,
312 "final_video_path": storyboard.final_video_path,
313 "total_duration": storyboard.total_duration,
314 "created_at": storyboard.created_at.isoformat() if storyboard.created_at else None,
315 "completed_at": storyboard.completed_at.isoformat() if storyboard.completed_at else None,
316 }
317
318 def _dict_to_storyboard(self, data: Dict[str, Any]) -> Storyboard:
319 """Convert dict to Storyboard instance"""
320 return Storyboard(
321 title=data["title"],
322 config=self._dict_to_config(data["config"]),
323 frames=[self._dict_to_frame(frame_data) for frame_data in data["frames"]],
324 content_metadata=self._dict_to_content_metadata(data["content_metadata"]) if data.get("content_metadata") else None,
325 final_video_path=data.get("final_video_path"),
326 total_duration=data.get("total_duration", 0.0),
327 created_at=datetime.fromisoformat(data["created_at"]) if data.get("created_at") else None,
328 completed_at=datetime.fromisoformat(data["completed_at"]) if data.get("completed_at") else None,
329 )
330
331 def _config_to_dict(self, config: StoryboardConfig) -> Dict[str, Any]:
332 """Convert StoryboardConfig to dict"""
333 return {
334 "task_id": config.task_id,
335 "n_storyboard": config.n_storyboard,
336 "min_narration_words": config.min_narration_words,
337 "max_narration_words": config.max_narration_words,
338 "min_image_prompt_words": config.min_image_prompt_words,
339 "max_image_prompt_words": config.max_image_prompt_words,
340 "video_fps": config.video_fps,
341 "tts_inference_mode": config.tts_inference_mode,
342 "voice_id": config.voice_id,
343 "tts_workflow": config.tts_workflow,
344 "tts_speed": config.tts_speed,
345 "ref_audio": config.ref_audio,
346 "media_width": config.media_width,
347 "media_height": config.media_height,
348 "media_workflow": config.media_workflow,
349 "frame_template": config.frame_template,
350 "template_params": config.template_params,
351 }
352
353 def _dict_to_config(self, data: Dict[str, Any]) -> StoryboardConfig:
354 """Convert dict to StoryboardConfig"""
355 return StoryboardConfig(
356 task_id=data.get("task_id"),
357 n_storyboard=data.get("n_storyboard", 5),
358 min_narration_words=data.get("min_narration_words", 5),
359 max_narration_words=data.get("max_narration_words", 20),
360 min_image_prompt_words=data.get("min_image_prompt_words", 30),
361 max_image_prompt_words=data.get("max_image_prompt_words", 60),
362 video_fps=data.get("video_fps", 30),
363 tts_inference_mode=data.get("tts_inference_mode", "local"),
364 voice_id=data.get("voice_id"),
365 tts_workflow=data.get("tts_workflow"),
366 tts_speed=data.get("tts_speed"),
367 ref_audio=data.get("ref_audio"),
368 media_width=data.get("media_width", data.get("image_width", 1024)), # Backward compatibility
369 media_height=data.get("media_height", data.get("image_height", 1024)), # Backward compatibility
370 media_workflow=data.get("media_workflow", data.get("image_workflow")), # Backward compatibility
371 frame_template=data.get("frame_template", "1080x1920/default.html"),
372 template_params=data.get("template_params"),
373 )
374
375 def _frame_to_dict(self, frame: StoryboardFrame) -> Dict[str, Any]:
376 """Convert StoryboardFrame to dict"""
377 return {
378 "index": frame.index,
379 "narration": frame.narration,
380 "image_prompt": frame.image_prompt,
381 "audio_path": frame.audio_path,
382 "media_type": frame.media_type,
383 "image_path": frame.image_path,
384 "video_path": frame.video_path,
385 "composed_image_path": frame.composed_image_path,
386 "video_segment_path": frame.video_segment_path,
387 "duration": frame.duration,
388 "created_at": frame.created_at.isoformat() if frame.created_at else None,
389 }
390
391 def _dict_to_frame(self, data: Dict[str, Any]) -> StoryboardFrame:
392 """Convert dict to StoryboardFrame"""
393 return StoryboardFrame(
394 index=data["index"],
395 narration=data["narration"],
396 image_prompt=data["image_prompt"],
397 audio_path=data.get("audio_path"),
398 media_type=data.get("media_type"),
399 image_path=data.get("image_path"),
400 video_path=data.get("video_path"),
401 composed_image_path=data.get("composed_image_path"),
402 video_segment_path=data.get("video_segment_path"),
403 duration=data.get("duration", 0.0),
404 created_at=datetime.fromisoformat(data["created_at"]) if data.get("created_at") else None,
405 )
406
407 def _content_metadata_to_dict(self, metadata: ContentMetadata) -> Dict[str, Any]:
408 """Convert ContentMetadata to dict"""
409 return {
410 "title": metadata.title,
411 "author": metadata.author,
412 "subtitle": metadata.subtitle,
413 "genre": metadata.genre,
414 "summary": metadata.summary,
415 "publication_year": metadata.publication_year,
416 "cover_url": metadata.cover_url,
417 }
418
419 def _dict_to_content_metadata(self, data: Dict[str, Any]) -> ContentMetadata:
420 """Convert dict to ContentMetadata"""
421 return ContentMetadata(
422 title=data["title"],
423 author=data.get("author"),
424 subtitle=data.get("subtitle"),
425 genre=data.get("genre"),
426 summary=data.get("summary"),
427 publication_year=data.get("publication_year"),
428 cover_url=data.get("cover_url"),
429 )
430
431 # ========================================================================
432 # Index Management (for fast listing)
433 # ========================================================================
434
435 def _ensure_index(self):
436 """Ensure index file exists, create if not"""
437 if not self.index_file.exists():
438 self._save_index({"version": "1.0", "tasks": []})
439
440 def _load_index(self) -> Dict[str, Any]:
441 """Load index from file"""
442 try:
443 with open(self.index_file, "r", encoding="utf-8") as f:
444 return json.load(f)
445 except Exception as e:
446 logger.error(f"Failed to load index: {e}")
447 return {"version": "1.0", "tasks": []}
448
449 def _save_index(self, index_data: Dict[str, Any]):
450 """Save index to file"""
451 try:
452 index_data["last_updated"] = datetime.now().isoformat()
453 with open(self.index_file, "w", encoding="utf-8") as f:
454 json.dump(index_data, f, ensure_ascii=False, indent=2)
455 except Exception as e:
456 logger.error(f"Failed to save index: {e}")
457
458 async def _update_index_for_task(self, task_id: str, metadata: Dict[str, Any]):
459 """Update index entry for a specific task"""
460 index = self._load_index()
461
462 # Try to get title from multiple sources
463 title = metadata.get("input", {}).get("title")
464 if not title or title == "":
465 # Try to get title from storyboard if input title is empty
466 storyboard = await self.load_storyboard(task_id)
467 if storyboard and storyboard.title:
468 title = storyboard.title
469 else:
470 # Fall back to using input text preview
471 input_text = metadata.get("input", {}).get("text", "")
472 if input_text:
473 # Use first 30 characters of input text as title
474 title = input_text[:30] + ("..." if len(input_text) > 30 else "")
475 else:
476 title = "Untitled"
477
478 # Extract key info for index
479 index_entry = {
480 "task_id": task_id,
481 "created_at": metadata.get("created_at"),
482 "completed_at": metadata.get("completed_at"),
483 "status": metadata.get("status", "unknown"),
484 "title": title,
485 "duration": metadata.get("result", {}).get("duration", 0),
486 "n_frames": metadata.get("result", {}).get("n_frames", 0),
487 "file_size": metadata.get("result", {}).get("file_size", 0),
488 "video_path": metadata.get("result", {}).get("video_path"),
489 }
490
491 # Update or append
492 tasks = index.get("tasks", [])
493 existing_idx = next((i for i, t in enumerate(tasks) if t["task_id"] == task_id), None)
494
495 if existing_idx is not None:
496 tasks[existing_idx] = index_entry
497 else:
498 tasks.append(index_entry)
499
500 index["tasks"] = tasks
501 self._save_index(index)
502
503 async def rebuild_index(self):
504 """Rebuild index by scanning all task directories"""
505 logger.info("Rebuilding task index...")
506 index = {"version": "1.0", "tasks": []}
507
508 # Scan all directories
509 for task_dir in self.output_dir.iterdir():
510 if not task_dir.is_dir() or task_dir.name.startswith("."):
511 continue
512
513 task_id = task_dir.name
514 metadata = await self.load_task_metadata(task_id)
515
516 if metadata:
517 # Try to get title from multiple sources
518 title = metadata.get("input", {}).get("title")
519 if not title or title == "":
520 # Try to get title from storyboard if input title is empty
521 storyboard = await self.load_storyboard(task_id)
522 if storyboard and storyboard.title:
523 title = storyboard.title
524 else:
525 # Fall back to using input text preview
526 input_text = metadata.get("input", {}).get("text", "")
527 if input_text:
528 # Use first 30 characters of input text as title
529 title = input_text[:30] + ("..." if len(input_text) > 30 else "")
530 else:
531 title = "Untitled"
532
533 # Add to index
534 index["tasks"].append({
535 "task_id": task_id,
536 "created_at": metadata.get("created_at"),
537 "completed_at": metadata.get("completed_at"),
538 "status": metadata.get("status", "unknown"),
539 "title": title,
540 "duration": metadata.get("result", {}).get("duration", 0),
541 "n_frames": metadata.get("result", {}).get("n_frames", 0),
542 "file_size": metadata.get("result", {}).get("file_size", 0),
543 "video_path": metadata.get("result", {}).get("video_path"),
544 })
545
546 self._save_index(index)
547 logger.info(f"Index rebuilt: {len(index['tasks'])} tasks")
548
549 # ========================================================================
550 # Paginated Listing
551 # ========================================================================
552
553 async def list_tasks_paginated(
554 self,
555 page: int = 1,
556 page_size: int = 20,
557 status: Optional[str] = None,
558 sort_by: str = "created_at",
559 sort_order: str = "desc"
560 ) -> Dict[str, Any]:
561 """
562 List tasks with pagination
563
564 Args:
565 page: Page number (1-indexed)
566 page_size: Items per page
567 status: Filter by status (optional)
568 sort_by: Sort field (created_at, completed_at, title, duration)
569 sort_order: Sort order (asc, desc)
570
571 Returns:
572 {
573 "tasks": [...], # List of task summaries
574 "total": 100, # Total matching tasks
575 "page": 1, # Current page
576 "page_size": 20, # Items per page
577 "total_pages": 5 # Total pages
578 }
579 """
580 index = self._load_index()
581 tasks = index.get("tasks", [])
582
583 # Filter by status
584 if status:
585 tasks = [t for t in tasks if t.get("status") == status]
586
587 # Sort
588 reverse = (sort_order == "desc")
589 if sort_by in ["created_at", "completed_at"]:
590 tasks.sort(
591 key=lambda t: datetime.fromisoformat(t.get(sort_by, "1970-01-01T00:00:00")),
592 reverse=reverse
593 )
594 elif sort_by in ["title", "duration", "n_frames"]:
595 tasks.sort(key=lambda t: t.get(sort_by, ""), reverse=reverse)
596
597 # Paginate
598 total = len(tasks)
599 total_pages = (total + page_size - 1) // page_size
600 start_idx = (page - 1) * page_size
601 end_idx = start_idx + page_size
602 page_tasks = tasks[start_idx:end_idx]
603
604 return {
605 "tasks": page_tasks,
606 "total": total,
607 "page": page,
608 "page_size": page_size,
609 "total_pages": total_pages,
610 }
611
612 # ========================================================================
613 # Statistics
614 # ========================================================================
615
616 async def get_statistics(self) -> Dict[str, Any]:
617 """
618 Get statistics about all tasks
619
620 Returns:
621 {
622 "total_tasks": 100,
623 "completed": 95,
624 "failed": 5,
625 "total_duration": 3600.5, # seconds
626 "total_size": 1024000000, # bytes
627 }
628 """
629 index = self._load_index()
630 tasks = index.get("tasks", [])
631
632 stats = {
633 "total_tasks": len(tasks),
634 "completed": len([t for t in tasks if t.get("status") == "completed"]),
635 "failed": len([t for t in tasks if t.get("status") == "failed"]),
636 "total_duration": sum(t.get("duration", 0) for t in tasks),
637 "total_size": sum(t.get("file_size", 0) for t in tasks),
638 }
639
640 return stats
641
642 # ========================================================================
643 # Delete Task
644 # ========================================================================
645
646 async def delete_task(self, task_id: str) -> bool:
647 """
648 Delete a task and all its files
649
650 Args:
651 task_id: Task ID to delete
652
653 Returns:
654 True if successful, False otherwise
655 """
656 try:
657 import shutil
658
659 task_dir = self.get_task_dir(task_id)
660 if task_dir.exists():
661 shutil.rmtree(task_dir)
662 logger.info(f"Deleted task directory: {task_dir}")
663
664 # Update index
665 index = self._load_index()
666 tasks = index.get("tasks", [])
667 tasks = [t for t in tasks if t["task_id"] != task_id]
668 index["tasks"] = tasks
669 self._save_index(index)
670
671 return True
672 except Exception as e:
673 logger.error(f"Failed to delete task {task_id}: {e}")
674 return False
675
676
676 lines PYTHON