| 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 | Template utility functions for size parsing and template management |
| 15 | """ |
| 16 | |
| 17 | import os |
| 18 | from pathlib import Path |
| 19 | from typing import List, Tuple, Optional, Literal |
| 20 | from pydantic import BaseModel, Field |
| 21 | import logging |
| 22 | |
| 23 | from pixelle_video.utils.os_util import ( |
| 24 | get_resource_path, |
| 25 | list_resource_files, |
| 26 | list_resource_dirs, |
| 27 | resource_exists |
| 28 | ) |
| 29 | |
| 30 | logger = logging.getLogger(__name__) |
| 31 | |
| 32 | |
| 33 | def parse_template_size(template_path: str) -> Tuple[int, int]: |
| 34 | """ |
| 35 | Parse video size from template path |
| 36 | |
| 37 | Args: |
| 38 | template_path: Template path like "templates/1080x1920/default.html" |
| 39 | or "1080x1920/default.html" |
| 40 | |
| 41 | Returns: |
| 42 | Tuple of (width, height) in pixels |
| 43 | |
| 44 | Raises: |
| 45 | ValueError: If template path format is invalid |
| 46 | |
| 47 | Examples: |
| 48 | >>> parse_template_size("templates/1080x1920/default.html") |
| 49 | (1080, 1920) |
| 50 | >>> parse_template_size("1920x1080/modern.html") |
| 51 | (1920, 1080) |
| 52 | """ |
| 53 | path = Path(template_path) |
| 54 | |
| 55 | # Get parent directory name (should be like "1080x1920") |
| 56 | dir_name = path.parent.name |
| 57 | |
| 58 | # Special case: if parent is "templates", go up one more level |
| 59 | if dir_name == "templates": |
| 60 | # This shouldn't happen in new structure, but handle it |
| 61 | raise ValueError( |
| 62 | f"Invalid template path format: {template_path}. " |
| 63 | f"Expected format: 'WIDTHxHEIGHT/template.html' or 'templates/WIDTHxHEIGHT/template.html'" |
| 64 | ) |
| 65 | |
| 66 | # Parse size from directory name |
| 67 | if 'x' not in dir_name: |
| 68 | raise ValueError( |
| 69 | f"Invalid size format in path: {template_path}. " |
| 70 | f"Directory name should be 'WIDTHxHEIGHT' (e.g., '1080x1920')" |
| 71 | ) |
| 72 | |
| 73 | try: |
| 74 | width_str, height_str = dir_name.split('x') |
| 75 | width = int(width_str) |
| 76 | height = int(height_str) |
| 77 | |
| 78 | # Sanity check |
| 79 | if width < 100 or height < 100 or width > 10000 or height > 10000: |
| 80 | raise ValueError(f"Invalid size dimensions: {width}x{height}") |
| 81 | |
| 82 | return (width, height) |
| 83 | except ValueError as e: |
| 84 | raise ValueError( |
| 85 | f"Failed to parse size from path: {template_path}. " |
| 86 | f"Expected format: 'WIDTHxHEIGHT/template.html' (e.g., '1080x1920/default.html'). " |
| 87 | f"Error: {e}" |
| 88 | ) |
| 89 | |
| 90 | |
| 91 | def list_available_sizes() -> List[str]: |
| 92 | """ |
| 93 | List all available video sizes (merged from templates/ and data/templates/) |
| 94 | |
| 95 | Returns: |
| 96 | List of size strings like ["1080x1920", "1920x1080", "1080x1080"] |
| 97 | |
| 98 | Examples: |
| 99 | >>> list_available_sizes() |
| 100 | ['1080x1920', '1920x1080', '1080x1080'] |
| 101 | """ |
| 102 | # Use new resource API to merge default and custom directories |
| 103 | all_dirs = list_resource_dirs("templates") |
| 104 | |
| 105 | # Filter to only valid size formats (WIDTHxHEIGHT) |
| 106 | sizes = [] |
| 107 | for dir_name in all_dirs: |
| 108 | if 'x' in dir_name: |
| 109 | try: |
| 110 | width, height = dir_name.split('x') |
| 111 | int(width) |
| 112 | int(height) |
| 113 | sizes.append(dir_name) |
| 114 | except (ValueError, AttributeError): |
| 115 | # Skip invalid directories |
| 116 | continue |
| 117 | |
| 118 | return sorted(sizes) |
| 119 | |
| 120 | |
| 121 | def list_templates_for_size(size: str) -> List[str]: |
| 122 | """ |
| 123 | List all templates available for a given size (merged from templates/ and data/templates/) |
| 124 | |
| 125 | Args: |
| 126 | size: Size string like "1080x1920" |
| 127 | |
| 128 | Returns: |
| 129 | List of template filenames (without path) like ["default.html", "modern.html"] |
| 130 | |
| 131 | Examples: |
| 132 | >>> list_templates_for_size("1080x1920") |
| 133 | ['cartoon.html', 'default.html', 'elegant.html', 'modern.html', ...] |
| 134 | """ |
| 135 | # Use new resource API to merge default and custom templates |
| 136 | all_files = list_resource_files("templates", size) |
| 137 | |
| 138 | # Filter to only HTML files |
| 139 | templates = [f for f in all_files if f.endswith('.html')] |
| 140 | |
| 141 | return sorted(templates) |
| 142 | |
| 143 | |
| 144 | def get_template_full_path(size: str, template_name: str) -> str: |
| 145 | """ |
| 146 | Get full template path from size and template name (checks data/templates/ first, then templates/) |
| 147 | |
| 148 | Args: |
| 149 | size: Size string like "1080x1920" |
| 150 | template_name: Template filename like "default.html" |
| 151 | |
| 152 | Returns: |
| 153 | Full path like "templates/1080x1920/default.html" or "data/templates/1080x1920/default.html" |
| 154 | |
| 155 | Raises: |
| 156 | FileNotFoundError: If template file doesn't exist in either location |
| 157 | |
| 158 | Examples: |
| 159 | >>> get_template_full_path("1080x1920", "default.html") |
| 160 | 'templates/1080x1920/default.html' |
| 161 | """ |
| 162 | # Use new resource API to search custom first, then default |
| 163 | try: |
| 164 | return get_resource_path("templates", size, template_name) |
| 165 | except FileNotFoundError: |
| 166 | available_templates = list_templates_for_size(size) |
| 167 | raise FileNotFoundError( |
| 168 | f"Template not found: {size}/{template_name}\n" |
| 169 | f"Available templates for size {size}: {available_templates}" |
| 170 | ) |
| 171 | |
| 172 | |
| 173 | class TemplateDisplayInfo(BaseModel): |
| 174 | """Template display information for UI layer""" |
| 175 | |
| 176 | name: str = Field(..., description="Template name without extension") |
| 177 | size: str = Field(..., description="Size string like '1080x1920'") |
| 178 | width: int = Field(..., description="Width in pixels") |
| 179 | height: int = Field(..., description="Height in pixels") |
| 180 | orientation: Literal['portrait', 'landscape', 'square'] = Field( |
| 181 | ..., |
| 182 | description="Video orientation" |
| 183 | ) |
| 184 | is_standard: bool = Field( |
| 185 | ..., |
| 186 | description="True only for standard sizes: 1080x1920, 1920x1080, 1080x1080" |
| 187 | ) |
| 188 | |
| 189 | |
| 190 | class TemplateInfo(BaseModel): |
| 191 | """Complete template information with path and display info""" |
| 192 | |
| 193 | template_path: str = Field(..., description="Full template path like '1080x1920/default.html'") |
| 194 | display_info: TemplateDisplayInfo = Field(..., description="Display information") |
| 195 | |
| 196 | |
| 197 | def format_template_display_info(template_name: str, size: str) -> TemplateDisplayInfo: |
| 198 | """ |
| 199 | Format template display information for UI |
| 200 | |
| 201 | Returns structured data for UI layer to handle display and i18n. |
| 202 | |
| 203 | Args: |
| 204 | template_name: Template filename like "default.html" |
| 205 | size: Size string like "1080x1920" |
| 206 | |
| 207 | Returns: |
| 208 | TemplateDisplayInfo object with name, size, dimensions, orientation, and standard flag |
| 209 | |
| 210 | Examples: |
| 211 | >>> info = format_template_display_info("default.html", "1080x1920") |
| 212 | >>> info.name |
| 213 | 'default' |
| 214 | >>> info.is_standard |
| 215 | True |
| 216 | |
| 217 | >>> info = format_template_display_info("custom.html", "1080x1921") |
| 218 | >>> info.orientation |
| 219 | 'portrait' |
| 220 | >>> info.is_standard |
| 221 | False |
| 222 | """ |
| 223 | # Keep full template name with .html extension |
| 224 | name = template_name |
| 225 | |
| 226 | # Parse size |
| 227 | width, height = map(int, size.split('x')) |
| 228 | |
| 229 | # Detect orientation |
| 230 | if height > width: |
| 231 | orientation = 'portrait' |
| 232 | elif width > height: |
| 233 | orientation = 'landscape' |
| 234 | else: |
| 235 | orientation = 'square' |
| 236 | |
| 237 | # Check if it's a standard size (only these three) |
| 238 | is_standard = (width, height) in [(1080, 1920), (1920, 1080), (1080, 1080)] |
| 239 | |
| 240 | return TemplateDisplayInfo( |
| 241 | name=name, |
| 242 | size=size, |
| 243 | width=width, |
| 244 | height=height, |
| 245 | orientation=orientation, |
| 246 | is_standard=is_standard |
| 247 | ) |
| 248 | |
| 249 | |
| 250 | def get_all_templates_with_info() -> List[TemplateInfo]: |
| 251 | """ |
| 252 | Get all templates with their display information |
| 253 | |
| 254 | Returns: |
| 255 | List of TemplateInfo objects |
| 256 | |
| 257 | Example: |
| 258 | >>> templates = get_all_templates_with_info() |
| 259 | >>> for t in templates: |
| 260 | ... print(f"{t.display_info.name} - {t.display_info.orientation}") |
| 261 | ... print(f" Path: {t.template_path}") |
| 262 | ... print(f" Standard: {t.display_info.is_standard}") |
| 263 | """ |
| 264 | result = [] |
| 265 | sizes = list_available_sizes() |
| 266 | |
| 267 | for size in sizes: |
| 268 | templates = list_templates_for_size(size) |
| 269 | for template in templates: |
| 270 | display_info = format_template_display_info(template, size) |
| 271 | full_path = f"{size}/{template}" |
| 272 | result.append(TemplateInfo( |
| 273 | template_path=full_path, |
| 274 | display_info=display_info |
| 275 | )) |
| 276 | |
| 277 | return result |
| 278 | |
| 279 | |
| 280 | def get_templates_grouped_by_size() -> dict: |
| 281 | """ |
| 282 | Get templates grouped by size |
| 283 | |
| 284 | Returns: |
| 285 | Dict with size as key, list of TemplateInfo as value |
| 286 | Ordered by orientation priority: portrait > landscape > square |
| 287 | |
| 288 | Example: |
| 289 | >>> grouped = get_templates_grouped_by_size() |
| 290 | >>> for size, templates in grouped.items(): |
| 291 | ... print(f"Size: {size}") |
| 292 | ... for t in templates: |
| 293 | ... print(f" - {t.display_info.name}") |
| 294 | """ |
| 295 | from collections import defaultdict |
| 296 | |
| 297 | templates = get_all_templates_with_info() |
| 298 | grouped = defaultdict(list) |
| 299 | |
| 300 | for t in templates: |
| 301 | grouped[t.display_info.size].append(t) |
| 302 | |
| 303 | # Sort groups by orientation priority: portrait > landscape > square |
| 304 | orientation_priority = {'portrait': 0, 'landscape': 1, 'square': 2} |
| 305 | |
| 306 | sorted_grouped = {} |
| 307 | for size in sorted(grouped.keys(), key=lambda s: ( |
| 308 | orientation_priority.get(grouped[s][0].display_info.orientation, 3), |
| 309 | s |
| 310 | )): |
| 311 | sorted_grouped[size] = sorted(grouped[size], key=lambda t: t.display_info.name) |
| 312 | |
| 313 | return sorted_grouped |
| 314 | |
| 315 | |
| 316 | def resolve_template_path(template_input: Optional[str]) -> str: |
| 317 | """ |
| 318 | Resolve template input to full path with validation (checks data/templates/ first, then templates/) |
| 319 | |
| 320 | Args: |
| 321 | template_input: Can be: |
| 322 | - None: Use default "1080x1920/image_default.html" |
| 323 | - "template.html": Use default size + this template |
| 324 | - "1080x1920/template.html": Full relative path |
| 325 | - "templates/1080x1920/template.html": Absolute-ish path (legacy) |
| 326 | - "data/templates/1080x1920/template.html": Custom path (legacy) |
| 327 | |
| 328 | Returns: |
| 329 | Resolved full path (custom if exists, otherwise default) |
| 330 | |
| 331 | Raises: |
| 332 | FileNotFoundError: If template doesn't exist in either location |
| 333 | |
| 334 | Examples: |
| 335 | >>> resolve_template_path(None) |
| 336 | 'templates/1080x1920/image_default.html' |
| 337 | >>> resolve_template_path("image_modern.html") |
| 338 | 'templates/1080x1920/image_modern.html' |
| 339 | >>> resolve_template_path("1920x1080/image_default.html") |
| 340 | 'templates/1920x1080/image_default.html' |
| 341 | """ |
| 342 | # Default case |
| 343 | if template_input is None: |
| 344 | template_input = "1080x1920/image_default.html" |
| 345 | |
| 346 | # Parse input to extract size and template name |
| 347 | size = None |
| 348 | template_name = None |
| 349 | |
| 350 | # Handle different input formats |
| 351 | if template_input.startswith("templates/") or template_input.startswith("data/templates/"): |
| 352 | # Legacy full path format - extract size and name |
| 353 | parts = Path(template_input).parts |
| 354 | if len(parts) >= 3: |
| 355 | size = parts[-2] |
| 356 | template_name = parts[-1] |
| 357 | elif '/' in template_input and 'x' in template_input.split('/')[0]: |
| 358 | # "1080x1920/template.html" format |
| 359 | size, template_name = template_input.split('/', 1) |
| 360 | else: |
| 361 | # Just template name - use default size |
| 362 | size = "1080x1920" |
| 363 | template_name = template_input |
| 364 | |
| 365 | # Backward compatibility: migrate "default.html" to "image_default.html" |
| 366 | if template_name == "default.html": |
| 367 | migrated_name = "image_default.html" |
| 368 | try: |
| 369 | # Try migrated name first |
| 370 | path = get_resource_path("templates", size, migrated_name) |
| 371 | logger.info(f"Backward compatibility: migrated '{template_input}' to '{size}/{migrated_name}'") |
| 372 | return path |
| 373 | except FileNotFoundError: |
| 374 | # Fall through to try original name |
| 375 | logger.warning(f"Migrated template '{size}/{migrated_name}' not found, trying original name") |
| 376 | |
| 377 | # Use resource API to resolve path (custom > default) |
| 378 | try: |
| 379 | return get_resource_path("templates", size, template_name) |
| 380 | except FileNotFoundError: |
| 381 | available_sizes = list_available_sizes() |
| 382 | raise FileNotFoundError( |
| 383 | f"Template not found: {size}/{template_name}\n" |
| 384 | f"Available sizes: {available_sizes}\n" |
| 385 | f"Hint: Use format 'SIZExSIZE/template.html' (e.g., '1080x1920/image_default.html')" |
| 386 | ) |
| 387 | |
| 388 | |
| 389 | def get_template_type(template_name: str) -> Literal['static', 'image', 'video']: |
| 390 | """ |
| 391 | Detect template type from template filename |
| 392 | |
| 393 | Template naming convention: |
| 394 | - static_*.html: Static style templates (no AI-generated media) |
| 395 | - image_*.html: Templates requiring AI-generated images |
| 396 | - video_*.html: Templates requiring AI-generated videos |
| 397 | |
| 398 | Args: |
| 399 | template_name: Template filename like "image_default.html" or "video_simple.html" |
| 400 | |
| 401 | Returns: |
| 402 | Template type: 'static', 'image', or 'video' |
| 403 | |
| 404 | Examples: |
| 405 | >>> get_template_type("static_simple.html") |
| 406 | 'static' |
| 407 | >>> get_template_type("image_default.html") |
| 408 | 'image' |
| 409 | >>> get_template_type("video_simple.html") |
| 410 | 'video' |
| 411 | """ |
| 412 | name = Path(template_name).name |
| 413 | |
| 414 | if name.startswith("static_"): |
| 415 | return "static" |
| 416 | elif name.startswith("video_"): |
| 417 | return "video" |
| 418 | elif name.startswith("image_"): |
| 419 | return "image" |
| 420 | else: |
| 421 | # Fallback: try to detect from legacy names |
| 422 | logger.warning( |
| 423 | f"Template '{template_name}' doesn't follow naming convention (static_/image_/video_). " |
| 424 | f"Defaulting to 'image' type." |
| 425 | ) |
| 426 | return "image" |
| 427 | |
| 428 | |
| 429 | def filter_templates_by_type( |
| 430 | templates: List[TemplateInfo], |
| 431 | template_type: Literal['static', 'image', 'video'] |
| 432 | ) -> List[TemplateInfo]: |
| 433 | """ |
| 434 | Filter templates by type |
| 435 | |
| 436 | Args: |
| 437 | templates: List of TemplateInfo objects |
| 438 | template_type: Type to filter by ('static', 'image', or 'video') |
| 439 | |
| 440 | Returns: |
| 441 | Filtered list of TemplateInfo objects |
| 442 | |
| 443 | Examples: |
| 444 | >>> all_templates = get_all_templates_with_info() |
| 445 | >>> image_templates = filter_templates_by_type(all_templates, 'image') |
| 446 | >>> len(image_templates) > 0 |
| 447 | True |
| 448 | """ |
| 449 | filtered = [] |
| 450 | for t in templates: |
| 451 | template_name = t.display_info.name |
| 452 | if get_template_type(template_name) == template_type: |
| 453 | filtered.append(t) |
| 454 | return filtered |
| 455 | |
| 456 | |
| 457 | def get_templates_grouped_by_size_and_type( |
| 458 | template_type: Optional[Literal['static', 'image', 'video']] = None |
| 459 | ) -> dict: |
| 460 | """ |
| 461 | Get templates grouped by size, optionally filtered by type |
| 462 | |
| 463 | Args: |
| 464 | template_type: Optional type filter ('static', 'image', or 'video') |
| 465 | |
| 466 | Returns: |
| 467 | Dict with size as key, list of TemplateInfo as value |
| 468 | Ordered by orientation priority: portrait > landscape > square |
| 469 | |
| 470 | Examples: |
| 471 | >>> # Get all templates |
| 472 | >>> all_grouped = get_templates_grouped_by_size_and_type() |
| 473 | |
| 474 | >>> # Get only image templates |
| 475 | >>> image_grouped = get_templates_grouped_by_size_and_type('image') |
| 476 | """ |
| 477 | from collections import defaultdict |
| 478 | |
| 479 | templates = get_all_templates_with_info() |
| 480 | |
| 481 | # Filter by type if specified |
| 482 | if template_type is not None: |
| 483 | templates = filter_templates_by_type(templates, template_type) |
| 484 | |
| 485 | grouped = defaultdict(list) |
| 486 | |
| 487 | for t in templates: |
| 488 | grouped[t.display_info.size].append(t) |
| 489 | |
| 490 | # Sort groups by orientation priority: portrait > landscape > square |
| 491 | orientation_priority = {'portrait': 0, 'landscape': 1, 'square': 2} |
| 492 | |
| 493 | sorted_grouped = {} |
| 494 | for size in sorted(grouped.keys(), key=lambda s: ( |
| 495 | orientation_priority.get(grouped[s][0].display_info.orientation, 3), |
| 496 | s |
| 497 | )): |
| 498 | sorted_grouped[size] = sorted(grouped[size], key=lambda t: t.display_info.name) |
| 499 | |
| 500 | return sorted_grouped |
| 501 | |
| 502 |