116 lines
3.1 KiB
Python
116 lines
3.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Experimental Apple MLX runtime helpers.
|
|
|
|
This package is intentionally small for now. It exists to grow the Apple-native
|
|
FastWan path in measurable steps: shape planning, primitive benchmarks, then
|
|
Wan block parity, then full DiT/runtime support.
|
|
"""
|
|
|
|
from fastvideo.mlx_runtime.fastwan import (
|
|
FastWanShape,
|
|
MLXQuantizationSpec,
|
|
MLXWanDiT,
|
|
MLXWanTransformerBlock,
|
|
UnsupportedMLXQuantizationError,
|
|
ensure_quantization_supported,
|
|
fastwan_shape,
|
|
fastwan_shape_from_config,
|
|
mlx_dit_from_diffusers_safetensors,
|
|
mlx_block_weights_from_torch,
|
|
mlx_block_weights_from_diffusers_safetensors,
|
|
quantization_support_error,
|
|
)
|
|
from fastvideo.mlx_runtime.checkpoint import (
|
|
load_mlx_dit_checkpoint,
|
|
save_mlx_dit_checkpoint,
|
|
)
|
|
from fastvideo.mlx_runtime.memory import (
|
|
AppliedMemoryLimits,
|
|
add_memory_limit_args,
|
|
apply_memory_limits,
|
|
gib_to_bytes,
|
|
)
|
|
from fastvideo.mlx_runtime.refine import (
|
|
DEFAULT_REFINE_SIGMA,
|
|
RefinePlan,
|
|
TwoPassResult,
|
|
default_refine_timesteps,
|
|
plan_refine_resolutions,
|
|
prepare_refine_latents,
|
|
refine_sigma_from_schedule,
|
|
run_dmd_loop,
|
|
run_two_pass_dmd,
|
|
upsample_latents_spatial,
|
|
)
|
|
from fastvideo.mlx_runtime.frame_upsample import (
|
|
DEFAULT_PIXEL_UPSAMPLE_MODE,
|
|
PIXEL_UPSAMPLE_MODES,
|
|
unsharp,
|
|
upsample_frame,
|
|
upsample_frames,
|
|
)
|
|
from fastvideo.mlx_runtime.fast_spatial import (
|
|
DEFAULT_FAST_SPATIAL_SHARPEN,
|
|
FastSpatialPlan,
|
|
apply_fast_spatial_upsample,
|
|
plan_fast_spatial,
|
|
resolve_spatial_mode,
|
|
)
|
|
from fastvideo.mlx_runtime.prompt_enhance import (
|
|
DEFAULT_ENHANCE_SYSTEM_PROMPT,
|
|
DEFAULT_MLX_LM_MODEL,
|
|
EnhanceResult,
|
|
enhance_prompt,
|
|
enhance_prompt_template,
|
|
enhance_result_as_metrics,
|
|
load_or_enhance_prompt,
|
|
)
|
|
|
|
__all__ = [
|
|
"AppliedMemoryLimits",
|
|
"DEFAULT_ENHANCE_SYSTEM_PROMPT",
|
|
"DEFAULT_MLX_LM_MODEL",
|
|
"DEFAULT_FAST_SPATIAL_SHARPEN",
|
|
"DEFAULT_PIXEL_UPSAMPLE_MODE",
|
|
"DEFAULT_REFINE_SIGMA",
|
|
"EnhanceResult",
|
|
"FastSpatialPlan",
|
|
"FastWanShape",
|
|
"MLXQuantizationSpec",
|
|
"MLXWanDiT",
|
|
"MLXWanTransformerBlock",
|
|
"RefinePlan",
|
|
"TwoPassResult",
|
|
"UnsupportedMLXQuantizationError",
|
|
"add_memory_limit_args",
|
|
"apply_fast_spatial_upsample",
|
|
"apply_memory_limits",
|
|
"enhance_prompt",
|
|
"enhance_prompt_template",
|
|
"enhance_result_as_metrics",
|
|
"ensure_quantization_supported",
|
|
"fastwan_shape",
|
|
"fastwan_shape_from_config",
|
|
"gib_to_bytes",
|
|
"load_mlx_dit_checkpoint",
|
|
"load_or_enhance_prompt",
|
|
"mlx_dit_from_diffusers_safetensors",
|
|
"mlx_block_weights_from_diffusers_safetensors",
|
|
"mlx_block_weights_from_torch",
|
|
"PIXEL_UPSAMPLE_MODES",
|
|
"default_refine_timesteps",
|
|
"plan_fast_spatial",
|
|
"plan_refine_resolutions",
|
|
"prepare_refine_latents",
|
|
"quantization_support_error",
|
|
"refine_sigma_from_schedule",
|
|
"resolve_spatial_mode",
|
|
"run_dmd_loop",
|
|
"run_two_pass_dmd",
|
|
"save_mlx_dit_checkpoint",
|
|
"unsharp",
|
|
"upsample_frame",
|
|
"upsample_frames",
|
|
"upsample_latents_spatial",
|
|
]
|