Files
hao-ai-lab-FastVideo/docs-site/scripts/prepare-python-docs.py
T

292 lines
13 KiB
Python

"""Build source-only Python reference pages with Griffe, without importing FastVideo."""
from __future__ import annotations
import html
import ast
import json
import re
import runpy
import shutil
import shlex
from pathlib import Path
from urllib.parse import urljoin, urlsplit, urlunsplit
import griffe
import markdown
import yaml
ROOT = Path(__file__).resolve().parents[2]
SITE = ROOT / "docs-site"
GENERATED = SITE / ".generated"
def source_controls(source: str) -> list[dict]:
"""Expose literal settings in maintained examples without importing them."""
labels = {
"prompt": "Prompt", "num_gpus": "GPU count", "height": "Height", "width": "Width",
"num_frames": "Frames", "seed": "Seed", "num_inference_steps": "Inference steps",
"guidance_scale": "Guidance scale", "fps": "FPS", "audio_end_in_s": "Audio duration",
"text_encoder_cpu_offload": "Offload text encoder", "vae_cpu_offload": "Offload VAE",
"dit_cpu_offload": "Offload transformer", "inference_torch_compile": "Compile transformer",
"video_decode_backend": "Video decoder",
}
lines = source.splitlines(keepends=True)
offsets = [0]
for line in lines:
offsets.append(offsets[-1] + len(line))
def position(line, column):
prefix = source[:offsets[line - 1]] + lines[line - 1].encode()[:column].decode()
return len(prefix.encode("utf-16-le")) // 2
settings = {}
def add(name, value):
if name not in labels or not isinstance(value, ast.Constant) or type(value.value) not in (str, int, float, bool):
return
key = json.dumps([name, value.value])
entry = settings.setdefault(key, {"name": name, "label": labels[name], "value": value.value,
"integer": type(value.value) is int, "spans": []})
entry["spans"].append([position(value.lineno, value.col_offset), position(value.end_lineno, value.end_col_offset)])
for node in ast.walk(ast.parse(source)):
if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
if node.func.attr in {"from_pretrained", "generate_video"}:
for keyword in node.keywords:
add(keyword.arg, keyword.value)
elif node.func.attr == "add_argument":
flags = [arg.value for arg in node.args if isinstance(arg, ast.Constant) and isinstance(arg.value, str)]
flag = next((flag for flag in flags if flag.startswith("--")), "")
for keyword in node.keywords:
if keyword.arg == "default":
add(flag.removeprefix("--").replace("-", "_"), keyword.value)
key = json.dumps([flag.removeprefix("--").replace("-", "_"), getattr(keyword.value, "value", None)])
if key in settings:
settings[key]["flag"] = flag
if any(item.arg == "type" and isinstance(item.value, ast.Name) and item.value.id == "float"
for item in node.keywords):
settings[key]["integer"] = False
settings[key]["booleanOptional"] = any(
item.arg == "action" and isinstance(item.value, ast.Attribute) and item.value.attr == "BooleanOptionalAction"
for item in node.keywords)
elif isinstance(node, ast.Assign):
for target in node.targets:
if isinstance(target, ast.Name) and target.id.isupper():
add(target.id.lower(), node.value)
result = list(settings.values())
seen = {}
for field in result:
seen[field["label"]] = seen.get(field["label"], 0) + 1
if seen[field["label"]] > 1:
field["label"] += f' ({seen[field["label"]]})'
return result
def signature(obj) -> str:
"""Preserve parameter kinds, annotations, defaults, and return annotations."""
if obj.is_attribute:
annotation = f": {obj.annotation}" if obj.annotation else ""
value = f" = {obj.value}" if obj.value is not None else ""
return f"{obj.name}{annotation}{value}"
parameters = list(getattr(obj, "parameters", []))
parts = []
keyword_separator = False
for i, param in enumerate(parameters):
if param.name in {"self", "cls"}:
continue
kind = param.kind.value
if kind == "keyword-only" and not keyword_separator:
parts.append("*")
keyword_separator = True
prefix = "**" if kind == "variadic keyword" else "*" if kind == "variadic positional" else ""
if prefix == "*":
keyword_separator = True
text = prefix + param.name
if param.annotation is not None:
text += f": {param.annotation}"
if param.default is not None and not prefix:
text += f" = {param.default}"
parts.append(text)
if kind == "positional-only" and (i + 1 == len(parameters) or parameters[i + 1].kind.value != kind):
parts.append("/")
result = f"{obj.name}({', '.join(parts)})"
if obj.is_function and obj.returns is not None:
result += f" -> {obj.returns}"
return result
def code_html(text: str) -> str:
return html.escape(text).replace("[", "[").replace("]", "]")
def prose(text: str) -> str:
# Docstrings are trusted repository source, rendered as Markdown rather than executed.
return markdown.markdown(text, extensions=["fenced_code", "tables", "sane_lists"])
def api_overview(text: str, package) -> str:
def reference(match):
label, target = match.groups()
path = (target or label).strip("`")
obj = package if path == package.path else package[path.removeprefix(package.path + ".")]
module = obj
while not module.is_module:
module = module.parent
route = "/api/" + module.path.replace(".", "/")
return f"[{label}]({route}#{path})"
text = re.sub(r"\[([^\]\n]+)\]\[([^\]\n]*)\]", reference, text)
def guide(match):
label, target = match.groups()
url = urlsplit(target)
if url.scheme or url.netloc or not re.search(r"\.mdx?$", url.path):
return match[0]
url = urlsplit(urljoin("/api/index.md", target))
route = re.sub(r"\.mdx?$", "", url.path)
route = re.sub(r"/index$", "", route) or "/"
return f"[{label}]({urlunsplit(('', '', route, url.query, url.fragment))})"
return prose(re.sub(r"\[([^\]\n]+)\]\(([^\s)]+)\)", guide, text))
def docstring(obj, level: int) -> str:
if not obj.docstring:
return ""
output = []
for section in obj.docstring.parsed:
kind = section.kind.value
if kind == "text":
output.append(prose(section.value))
continue
title = section.title or kind.replace("_", " ").capitalize()
output.append(f"<h{level}>{html.escape(title)}</h{level}>")
if kind == "examples":
for example_kind, value in section.value:
output.append(f"<pre><code>{html.escape(value)}</code></pre>" if example_kind.value == "examples" else prose(value))
elif isinstance(section.value, list):
output.append("<dl>")
for item in section.value:
name = getattr(item, "name", "") or ""
annotation = getattr(item, "annotation", "") or ""
label = f"{name}: {annotation}" if name and annotation else str(name or annotation)
output.append(f"<dt><code>{html.escape(label)}</code></dt><dd>{prose(getattr(item, 'description', ''))}</dd>")
output.append("</dl>")
else:
output.append(prose(getattr(section.value, "description", str(section.value))))
return "\n".join(output)
def symbol(obj, path: str, level: int = 2, members: list[str] | None = None, source: bool = True) -> str:
level = min(level, 6)
output = [f'<h{level} id="{html.escape(path)}"><code>{html.escape(path)}</code></h{level}>']
output.append(f'<pre class="doc-signature"><code>{code_html(signature(obj))}</code></pre>')
output.append(docstring(obj, min(level + 1, 6)))
if obj.is_class:
constructor = obj.members.get("__init__")
if constructor is not None:
output.append(docstring(constructor, min(level + 1, 6)))
if source and not obj.is_alias:
try:
code = obj.source
except (ValueError, KeyError):
code = ""
if code:
output.append(f'<details><summary>Source</summary><pre><code>{code_html(code)}</code></pre></details>')
if obj.is_class:
for name, member in obj.members.items():
if members is not None and name not in members:
continue
if name.startswith("_") or member.is_module:
continue
if member.is_attribute and not member.docstring:
continue
try:
output.append(symbol(member, f"{path}.{name}", min(level + 1, 6), source=source))
except griffe.AliasResolutionError:
continue
return "\n".join(output)
def write_page(route: str, title: str, body: str, anchors: dict) -> None:
# Preserve incoming fragments, including historical category labels whose layout changed.
present = set(re.findall(r'id="([^"]+)"', body))
aliases = "\n".join(f'<a id="{html.escape(h["id"])}"></a>' for h in anchors.get(route, []) if h["id"] not in present)
target = GENERATED / "python-reference" / route / "index.md"
target.parent.mkdir(parents=True, exist_ok=True)
target.write_text("---\n" + yaml.safe_dump({"title": title}, sort_keys=False) +
f'---\n\n<div class="python-reference">\n{aliases}\n{body}\n</div>\n')
def prepare() -> None:
GENERATED.mkdir(exist_ok=True)
examples = runpy.run_path(str(ROOT / "docs/generate_examples.py"))
examples["validate_cookbook"]()
examples["generate_examples"](generate_main_index=True)
source_dir = GENERATED / "source-examples"
shutil.rmtree(source_dir, ignore_errors=True)
source_dir.mkdir()
for recipe in json.loads((ROOT / "docs/assets/cookbook-recipes.json").read_text())["recipes"]:
path = ROOT / recipe["source"]
if path.suffix == ".py":
source = path.read_text()
controls = source_controls(source)
argv = shlex.split(recipe["command"])
for field in controls:
flag = field.get("flag")
positions = [i for i, argument in enumerate(argv[:-1]) if argument == flag]
if positions and type(field["value"]) is not bool:
converter = float if isinstance(field["value"], (int, float)) and not field["integer"] else type(field["value"])
field["value"] = converter(argv[positions[-1] + 1])
elif field.get("booleanOptional"):
if flag in argv:
field["value"] = True
elif flag.replace("--", "--no-", 1) in argv:
field["value"] = False
(source_dir / f'{recipe["id"]}.json').write_text(json.dumps({"source": source, "controls": controls}))
anchors = json.loads((SITE / "migration-anchors.json").read_text())
shutil.rmtree(GENERATED / "python-reference", ignore_errors=True)
package = griffe.load("fastvideo", search_paths=[ROOT], allow_inspection=False, resolve_aliases=True,
resolve_external=False, docstring_parser="google", docstring_options={"warnings": False})
modules = []
def visit(module):
if module.is_alias or any(part.startswith("_") or part in {"tests", "third_party"} for part in module.path.split(".")):
return
modules.append(module)
for child in module.modules.values():
visit(child)
visit(package)
for module in sorted(modules, key=lambda item: item.path):
body = [f'<a id="{module.path}"></a>', docstring(module, 2)]
for name, member in module.members.items():
if name.startswith("_"):
continue
try:
# External imports belong to their upstream reference, not FastVideo's API.
if member.is_alias and not member.target_path.startswith("fastvideo."):
continue
if member.is_module or member.is_attribute and not member.docstring:
continue
body.append(symbol(member, f"{module.path}.{name}"))
except griffe.AliasResolutionError:
continue
write_page("api/" + module.path.replace(".", "/"), module.name, "\n".join(body), anchors)
write_page("api", "Python API reference", api_overview((ROOT / "docs/api/index.md").read_text(), package), anchors)
v1 = (ROOT / "docs/getting_started/v1_api.md").read_text()
def reference(match):
path, level, members = match.groups()
return "\n" + symbol(package[path.removeprefix("fastvideo.")], path,
int(level), members.split(","), source=False) + "\n"
v1 = re.sub(r'<api-reference path="([^"]+)" heading="(\d+)" members="([^"]+)" />', reference, v1)
write_page("getting_started/v1_api", "V1 API", prose(v1.removeprefix("# V1 API\n")), anchors)
(GENERATED / "old-headings.json").write_text(json.dumps(anchors))
print(f"Rendered {len(modules) + 1} Python reference pages directly from source, without MkDocs or runtime imports.")
if __name__ == "__main__":
prepare()