nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
bc27a032c587de330fd29b70dea5c02eeae6c59f
hao-ai-lab-FastVideo/docs/generate_examples.py
T

567 lines
22 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/docs/source/generate_examples.py
import itertools
import os
import re
from dataclasses import dataclass, field
from pathlib import Path
ROOT_DIR = Path(__file__).parent.parent.resolve()
ROOT_DIR_RELATIVE = '../..'
EXAMPLE_DIR = ROOT_DIR / "examples"
EXAMPLE_DOC_DIR = ROOT_DIR / "docs/getting_started/examples"
GITHUB_REPO = "hao-ai-lab/FastVideo" # Update this to your repo
def fix_case(text: str) -> str:
subs = {
"api": "API",
"cli": "CLI",
"cpu": "CPU",
"llm": "LLM",
"tpu": "TPU",
"aqlm": "AQLM",
"gguf": "GGUF",
"lora": "LoRA",
"rlhf": "RLHF",
"vllm": "vLLM",
"openai": "OpenAI",
"multilora": "MultiLoRA",
"mlpspeculator": "MLPSpeculator",
"finetune": "Finetune",
"distillation": "Distillation",
"wan": "Wan",
"i2v": "I2V",
"t2v": "T2V",
"1.3b": "1.3B",
"14b": "14B",
"480p": "480P",
"720p": "720P",
r"fp\d+": lambda x: x.group(0).upper(), # e.g. fp16, fp32
r"int\d+": lambda x: x.group(0).upper(), # e.g. int8, int16
}
for pattern, repl in subs.items():
text = re.sub(rf'\b{pattern}\b', repl, text,
flags=re.IGNORECASE) # type: ignore[call-overload]
return text
@dataclass
class Index:
"""
Index class to generate a structured document index.
Attributes:
path (Path): The path save the index file to.
title (str): The title of the index.
description (str): A brief description of the index.
caption (str): An optional caption for the table of contents.
maxdepth (int): The maximum depth of the table of contents. Defaults to 1.
documents (list[str]): A list of document paths to include in the index. Defaults to an empty list.
Methods:
generate() -> str:
Generates the index content as a string in the specified format.
""" # noqa: E501
path: Path
title: str
description: str
caption: str
maxdepth: int = 1
documents: list[str] = field(default_factory=list)
def generate(self) -> str:
content = f"# {self.title}\n\n{self.description}\n\n"
if self.caption:
content += f"## {self.caption}\n\n"
# Generate a simple list of links for MkDocs
for doc in self.documents:
# Convert document path to proper link
doc_link = doc.replace("\\", "/")
# Get just the filename for the link text
doc_title = fix_case(Path(doc).stem.replace("_", " ").title())
content += f"- [{doc_title}]({doc_link}.md)\n"
content += "\n"
return content
@dataclass
class Example:
"""
Example class for generating documentation content from a given path.
Attributes:
path (Path): The path to the main directory or file.
category (str): The category of the document.
main_file (Path): The main file in the directory.
other_files (list[Path]): list of other files in the directory.
title (str): The title of the document.
Methods:
__post_init__(): Initializes the main_file, other_files, and title attributes.
determine_main_file() -> Path: Determines the main file in the given path.
determine_other_files() -> list[Path]: Determines other files in the directory excluding the main file.
determine_title() -> str: Determines the title of the document.
generate() -> str: Generates the documentation content.
""" # noqa: E501
path: Path
category: str | None = None
main_file: Path = field(init=False)
other_files: list[Path] = field(init=False)
title: str = field(init=False)
def __post_init__(self):
self.main_file = self.determine_main_file()
self.other_files = self.determine_other_files()
self.title = self.determine_title()
def determine_main_file(self) -> Path:
"""
Determines the main file in the given path.
If the path is a file, it returns the path itself. Otherwise, it searches
for Markdown files (*.md) in the directory and returns the first one found.
Returns:
Path: The main file path, either the original path if it's a file or the first
Markdown file found in the directory.
Raises:
IndexError: If no Markdown files are found in the directory.
""" # noqa: E501
if self.path.is_file():
return self.path
markdown_files = sorted(self.path.glob("*.md"))
if not markdown_files:
raise IndexError(f"No Markdown files found in {self.path}")
readme_files = [
f for f in markdown_files if f.name.lower() == "readme.md"
]
if readme_files:
return readme_files[0]
return markdown_files[0]
def determine_other_files(self) -> list[Path]:
"""
Determine other files in the directory excluding the main file.
This method checks if the given path is a file. If it is, it returns an empty list.
Otherwise, it recursively searches through the directory and returns a list of all
files that are not the main file.
Returns:
list[Path]: A list of Path objects representing the other files in the directory.
""" # noqa: E501
if self.path.is_file():
return []
is_other_file = lambda file: file.is_file() and file != self.main_file
return [file for file in self.path.rglob("*")
if is_other_file(file)] # type: ignore[no-untyped-call]
def determine_title(self) -> str:
return fix_case(self.path.stem.replace("_", " ").title())
def generate(self) -> str:
# Create GitHub link to source
github_path = str(self.path.relative_to(ROOT_DIR)).replace("\\", "/")
github_url = f"https://github.com/{GITHUB_REPO}/blob/main/{github_path}"
content = f"**Source:** [{github_path}]({github_url})\n\n"
# Add title for code files
if self.main_file.suffix != ".md":
content += f"# {self.title}\n\n"
# Include main file content
if self.main_file.suffix == ".md":
# For markdown files, include the content directly
with open(self.main_file, encoding='utf-8') as f:
content += f.read() + "\n\n"
else:
# For code files, use code blocks
language = self.main_file.suffix[1:] if self.main_file.suffix else ""
with open(self.main_file, encoding='utf-8') as f:
file_content = f.read()
content += f"```{language}\n{file_content}\n```\n\n"
if not self.other_files:
return content
content += "## Additional Files\n\n"
# Define binary/non-text file extensions to skip
binary_extensions = {
'.mp4', '.avi', '.mov', '.mkv', '.gif', '.jpg', '.jpeg', '.png',
'.webp', '.bmp', '.pdf', '.zip', '.tar', '.gz', '.mp3', '.wav'
}
for file in sorted(self.other_files):
# Skip binary files
if file.suffix.lower() in binary_extensions:
continue
file_rel_path = file.relative_to(self.path)
# Use collapsible admonition syntax for MkDocs
content += f"??? note \"{file_rel_path}\"\n\n"
try:
if file.suffix == ".md":
# Include markdown content with indentation
with open(file, encoding='utf-8') as f:
for line in f:
content += f" {line}"
else:
# Include code with proper formatting
language = file.suffix[1:] if file.suffix else ""
with open(file, encoding='utf-8') as f:
file_content = f.read()
# Indent the code block for the admonition
content += f" ```{language}\n"
for line in file_content.split('\n'):
content += f" {line}\n"
content += " ```\n"
content += "\n"
except UnicodeDecodeError:
# Skip files that can't be decoded as UTF-8
continue
return content
@dataclass
class NestedStructure:
"""Helper class to manage nested documentation structures for training/distillation."""
category: str
method: str
model: str
dataset: str
example: Example
@property
def filename(self) -> str:
return f"{self.model}_{self.dataset}"
@property
def title(self) -> str:
return fix_case(self.dataset.replace('_', ' '))
@property
def description(self) -> str:
category_name = self.category.title()
return f"{category_name} example using the {self.dataset} dataset with the {self.model} model."
def create_category_indices() -> dict[str, Index]:
"""Create category indices with their respective configurations."""
main_index_dir = ROOT_DIR / "docs/examples"
if not main_index_dir.exists():
main_index_dir.mkdir(parents=True)
category_indices = {
"inference":
Index(
path=ROOT_DIR /
"docs/inference/examples/examples_inference_index.md",
title="🚀 Examples",
description=
"Inference examples demonstrate how to use FastVideo inference. We recommend starting with [basic.md](basic.md).",
caption="Examples",
maxdepth=1,
),
"training":
Index(
path=ROOT_DIR / "docs/training/examples/examples_training_index.md",
title="🚀 Examples",
description=
"Training examples demonstrate how to use FastVideo training.",
caption="Examples",
maxdepth=3,
),
"distillation":
Index(
path=ROOT_DIR /
"docs/distillation/examples/examples_distillation_index.md",
title="🚀 Examples",
description=
"Distillation examples demonstrate how to use FastVideo distillation.",
caption="Examples",
maxdepth=3,
),
}
# Ensure all category doc directories exist
for index in category_indices.values():
if not index.path.parent.exists():
index.path.parent.mkdir(parents=True)
return category_indices
def find_examples(category_indices: dict[str, Index],
generate_main_index: bool) -> list[Example]:
"""Find all examples from the examples directory."""
examples = []
glob_patterns = ["*.py", "*.md", "*.sh"]
# Map category names to actual directory names
category_dir_mapping = {
"distillation": "distill", # examples/distill/ -> distillation category
}
# Find categorised examples
for category in category_indices:
# Use mapped directory name if available, otherwise use category name
dir_name = category_dir_mapping.get(category, category)
category_dir = EXAMPLE_DIR / dir_name
# Skip if directory doesn't exist
if not category_dir.exists():
continue
globs = [category_dir.glob(pattern) for pattern in glob_patterns]
for path in itertools.chain(*globs):
examples.append(Example(path, category))
# Find examples in subdirectories (recursively)
for path in category_dir.glob("**/*.md"):
examples.append(Example(path.parent, category))
# Find uncategorised examples only if we're generating a main index
if generate_main_index:
globs = [EXAMPLE_DIR.glob(pattern) for pattern in glob_patterns]
for path in itertools.chain(*globs):
examples.append(Example(path))
# Find examples in subdirectories
for path in EXAMPLE_DIR.glob("*/*.md"):
# Skip categorised examples
if path.parent.name in category_indices:
continue
examples.append(Example(path.parent))
return examples
def create_nested_structures(
examples: list[Example]
) -> dict[str, dict[str, dict[str, dict[str, NestedStructure]]]]:
"""Create nested structures for training and distillation categories."""
nested_structures: dict[str, dict[str, dict[str,
dict[str,
NestedStructure]]]] = {}
# Map category names to actual directory names
category_dir_mapping = {
"distillation": "distill",
}
for example in examples:
if example.category not in ["training", "distillation"]:
continue
# Use mapped directory name if available
dir_name = category_dir_mapping.get(example.category, example.category)
category_dir = EXAMPLE_DIR / dir_name
relative_path = example.path.relative_to(category_dir)
path_parts = relative_path.parts
if example.category == "training":
# For training examples like finetune/wan_i2v_14b_480p/crush_smol
if len(path_parts) >= 3:
method = path_parts[0] # e.g., "finetune"
model = path_parts[1] # e.g., "wan_i2v_14b_480p"
dataset = path_parts[2] # e.g., "crush_smol"
# Initialize nested structure
if example.category not in nested_structures:
nested_structures[example.category] = {}
if method not in nested_structures[example.category]:
nested_structures[example.category][method] = {}
if model not in nested_structures[example.category][method]:
nested_structures[example.category][method][model] = {}
# Store the nested structure
nested_structures[
example.category][method][model][dataset] = NestedStructure(
category=example.category,
method=method,
model=model,
dataset=dataset,
example=example)
elif example.category == "distillation" and len(path_parts) >= 2:
# For distillation examples like Wan2.1-T2V/Wan-Syn-Data-480P
model = path_parts[0] # e.g., "Wan2.1-T2V"
dataset = path_parts[1] # e.g., "Wan-Syn-Data-480P"
method = "DMD" # Default method for distillation
# Initialize nested structure
if example.category not in nested_structures:
nested_structures[example.category] = {}
if method not in nested_structures[example.category]:
nested_structures[example.category][method] = {}
if model not in nested_structures[example.category][method]:
nested_structures[example.category][method][model] = {}
# Store the nested structure
nested_structures[
example.category][method][model][dataset] = NestedStructure(
category=example.category,
method=method,
model=model,
dataset=dataset,
example=example)
return nested_structures
def generate_flat_examples(examples: list[Example],
category_indices: dict[str, Index],
examples_index: Index | None,
generate_main_index: bool) -> None:
"""Generate documentation for flat structure examples (inference, etc.)."""
for example in examples:
if example.category in ["training", "distillation"]:
continue # Skip nested structure examples
# Determine which index to use for this example
if example.category is not None and example.category in category_indices:
index = category_indices[example.category]
elif generate_main_index:
assert examples_index is not None
index = examples_index
else:
continue
# Generate the example documentation
doc_path = index.path.parent / f"{example.path.stem}.md"
with open(doc_path, "w+") as f:
f.write(example.generate())
index.documents.append(example.path.stem)
def generate_nested_examples(nested_structures: dict[str, dict[str, dict[
str, dict[str, NestedStructure]]]], category_indices: dict[str,
Index]) -> None:
"""Generate documentation for nested structure examples (training, distillation)."""
for category_name in ["training", "distillation"]:
if category_name not in category_indices or category_name not in nested_structures:
continue
category_index = category_indices[category_name]
category_base_dir = category_index.path.parent
for method, models in nested_structures[category_name].items():
# Create method-level index
method_index = Index(path=category_base_dir / f"{method}.md",
title=fix_case(method),
description=f"Examples using {method}.",
caption=f"{fix_case(method)} Examples",
maxdepth=2)
for model, datasets in models.items():
# Generate dataset examples using the Example class
for dataset, nested_struct in datasets.items():
doc_path = category_base_dir / f"{nested_struct.filename}.md"
with open(doc_path, "w+") as f:
f.write(nested_struct.example.generate())
# Create model-level index
model_index = Index(
path=category_base_dir / f"{model}.md",
title=fix_case(model.replace('_', ' ')),
description=f"Examples for the {model} model.",
caption=f"{fix_case(model.replace('_', ' '))} Datasets",
maxdepth=1)
# Add dataset indices to model index
for dataset, nested_struct in datasets.items():
model_index.documents.append(nested_struct.filename)
# Write model index
with open(model_index.path, "w+") as f:
f.write(model_index.generate())
# Add model to method index
method_index.documents.append(model)
# Write method index
with open(method_index.path, "w+") as f:
f.write(method_index.generate())
# Add method to main category index
category_index.documents.append(method)
def generate_examples(generate_main_index: bool = False) -> None:
"""
Generate example documentation.
Args:
generate_main_index (bool): Whether to generate the main examples index.
If False, only category-specific indices will be generated.
"""
# Create category indices
category_indices = create_category_indices()
# Create the main examples index only if requested
examples_index = None
if generate_main_index:
main_index_dir = ROOT_DIR / "docs/examples"
examples_index = Index(
path=main_index_dir / "examples_index.md",
title="💡 Examples",
description=
"A collection of examples demonstrating usage of FastVideo.\n\n"
f"All documented examples are autogenerated using [generate_examples.py](https://github.com/{GITHUB_REPO}/blob/main/docs/generate_examples.py) "
f"from examples found in the [examples](https://github.com/{GITHUB_REPO}/tree/main/examples) directory.",
caption="Examples",
maxdepth=2)
# Find all examples
examples = find_examples(category_indices, generate_main_index)
# Create nested structures for training and distillation
nested_structures = create_nested_structures(examples)
# Generate flat structure examples (inference, etc.)
generate_flat_examples(examples, category_indices, examples_index,
generate_main_index)
# Generate nested structure examples (training, distillation)
generate_nested_examples(nested_structures, category_indices)
# Generate the index files for categories
for category_index in category_indices.values():
if category_index.documents:
# Add to main index if it exists
if generate_main_index and examples_index:
main_index_dir = examples_index.path.parent
rel_path = os.path.relpath(category_index.path,
start=main_index_dir)
examples_index.documents.insert(
0,
str(rel_path).replace("\\", "/").replace(".md", ""))
# Write the category index file
with open(category_index.path, "w+") as f:
f.write(category_index.generate())
# Write the main index file if requested
if generate_main_index and examples_index:
with open(examples_index.path, "w+") as f:
f.write(examples_index.generate())
def on_pre_build_hook(config, **kwargs):
"""
MkDocs hook to generate examples before building the documentation.
This function is called automatically by the mkdocs-simple-hooks plugin.
"""
print("Generating example documentation...")
generate_examples(generate_main_index=True)
print("Example documentation generated successfully!")
if __name__ == "__main__":
print("Generating example documentation...")
generate_examples(generate_main_index=True)
print("Example documentation generated successfully!")
Reference in New Issue View Git Blame Copy Permalink