ac48bb3cd1
Co-authored-by: Aryan Kumar <aryank@Aryans-Mac-Studio.local> Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: CodeRabbit <noreply@coderabbit.ai>
131 lines
3.4 KiB
Python
131 lines
3.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Apple Silicon MLX inference helpers."""
|
|
|
|
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.checkpoint_compat import (
|
|
UnsupportedMLXCheckpointError,
|
|
discover_mlx_checkpoint,
|
|
raise_if_unsupported_mlx_checkpoint,
|
|
resolve_mlx_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,
|
|
)
|
|
from fastvideo.mlx_runtime.minimax_h3_pipeline import (
|
|
FastTemporalPlan,
|
|
GenerationResult,
|
|
MiniMaxH3MLXPipeline,
|
|
plan_fast_temporal,
|
|
)
|
|
|
|
__all__ = [
|
|
"AppliedMemoryLimits",
|
|
"DEFAULT_ENHANCE_SYSTEM_PROMPT",
|
|
"DEFAULT_MLX_LM_MODEL",
|
|
"DEFAULT_FAST_SPATIAL_SHARPEN",
|
|
"DEFAULT_PIXEL_UPSAMPLE_MODE",
|
|
"DEFAULT_REFINE_SIGMA",
|
|
"EnhanceResult",
|
|
"FastSpatialPlan",
|
|
"FastTemporalPlan",
|
|
"FastWanShape",
|
|
"GenerationResult",
|
|
"MLXQuantizationSpec",
|
|
"MLXWanDiT",
|
|
"MLXWanTransformerBlock",
|
|
"MiniMaxH3MLXPipeline",
|
|
"RefinePlan",
|
|
"TwoPassResult",
|
|
"UnsupportedMLXCheckpointError",
|
|
"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",
|
|
"discover_mlx_checkpoint",
|
|
"load_mlx_dit_checkpoint",
|
|
"load_or_enhance_prompt",
|
|
"raise_if_unsupported_mlx_checkpoint",
|
|
"resolve_mlx_checkpoint",
|
|
"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_fast_temporal",
|
|
"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",
|
|
]
|