3333import numpy as np
3434
3535from fastvideo .logger import init_logger
36+ from fastvideo .mlx_runtime .frame_upsample import (
37+ DEFAULT_PIXEL_UPSAMPLE_MODE ,
38+ PIXEL_UPSAMPLE_MODES ,
39+ upsample_frames ,
40+ )
3641from fastvideo .mlx_runtime .minimax_h3 import (
3742 H3_MANIFEST_FILENAME ,
3843 MINIMAX_H3_AUDIO_SHIFT ,
@@ -99,6 +104,72 @@ def plan_fast_temporal(target_frames: int, factor: int = 2) -> FastTemporalPlan:
99104 )
100105
101106
107+ # Resampling from a smaller decode softens output the same way on every
108+ # runtime; 0.4 matches the tuned Wan default without the halos that show
109+ # up by ~0.8.
110+ DEFAULT_FAST_SPATIAL_SHARPEN = 0.4
111+
112+
113+ @dataclass (frozen = True )
114+ class FastSpatialPlan :
115+ """Reduced-canvas geometry for spatial fast mode (RIFE's spatial twin)."""
116+
117+ target_height : int
118+ target_width : int
119+ stage1_height : int
120+ stage1_width : int
121+ canvas_height : int
122+ canvas_width : int
123+ scale : int
124+ upsample_mode : str
125+ sharpen : float
126+
127+
128+ def plan_fast_spatial (
129+ height : int ,
130+ width : int ,
131+ * ,
132+ scale : int = 2 ,
133+ upsample_mode : str = DEFAULT_PIXEL_UPSAMPLE_MODE ,
134+ sharpen : float = DEFAULT_FAST_SPATIAL_SHARPEN ,
135+ ) -> FastSpatialPlan :
136+ """Choose the smallest H3-valid canvas that covers ``target / scale``.
137+
138+ H3 geometry rounds *up* to the 32px model grid and center-crops after
139+ decode — the same convention plain 720p generation uses via
140+ ``_model_canvas_size`` — so no size the full-resolution path accepts is
141+ rejected here. The return trip to the target size runs in pixel space
142+ after the VAE decode, never on latents; see
143+ :mod:`fastvideo.mlx_runtime.frame_upsample` for why.
144+ """
145+ if scale < 2 :
146+ raise ValueError (f"fast-spatial scale must be at least 2, got { scale } )
147+ if upsample_mode not in PIXEL_UPSAMPLE_MODES :
148+ raise ValueError (f"Unsupported upsample mode: { upsample_mode !r}
149+ f"(expected one of { ', ' .join (PIXEL_UPSAMPLE_MODES )} )
150+ if sharpen < 0 :
151+ raise ValueError (f"fast_spatial_sharpen must be non-negative, got { sharpen } )
152+ target_canvas_height , target_canvas_width = _model_canvas_size (height , width )
153+ stage1_height = math .ceil (height / scale )
154+ stage1_width = math .ceil (width / scale )
155+ canvas_height , canvas_width = _model_canvas_size (stage1_height , stage1_width )
156+ if canvas_height * canvas_width >= target_canvas_height * target_canvas_width :
157+ raise ValueError (
158+ f"fast-spatial scale { scale } { height } { width }
159+ f"(stage-1 canvas { canvas_width } { canvas_height } { target_canvas_width } { target_canvas_height } )
160+ return FastSpatialPlan (
161+ target_height = height ,
162+ target_width = width ,
163+ stage1_height = stage1_height ,
164+ stage1_width = stage1_width ,
165+ canvas_height = canvas_height ,
166+ canvas_width = canvas_width ,
167+ scale = scale ,
168+ upsample_mode = upsample_mode ,
169+ sharpen = sharpen ,
170+ )
171+
172+
102173def _model_canvas_size (height : int , width : int ) -> tuple [int , int ]:
103174 """Round an exact output size up to H3's 32-pixel model grid."""
104175 if height <= 0 or width <= 0 :
@@ -202,10 +273,16 @@ def _validate_checkpoint_step_ladder(checkpoint_dir: str | Path, num_steps: int)
202273 f"{ num_steps } )
203274
204275
205- def _preflight_media_dependencies (* , fast : bool , fast_sharpen : float , rife_weights_dir : str | Path | None ) -> None :
276+ def _preflight_media_dependencies (* ,
277+ fast : bool ,
278+ fast_sharpen : float ,
279+ rife_weights_dir : str | Path | None ,
280+ fast_spatial : bool = False ) -> None :
206281 """Fail before conditioning when required output dependencies are unavailable."""
207282 if shutil .which ("ffmpeg" ) is None :
208283 raise RuntimeError ("ffmpeg is required for MP4 muxing; install it before generation." )
284+ if fast_spatial and importlib .util .find_spec ("cv2" ) is None :
285+ raise RuntimeError ("OpenCV is required for --fast-spatial resampling." )
209286 if not fast :
210287 return
211288 if fast_sharpen > 0 and importlib .util .find_spec ("cv2" ) is None :
@@ -629,6 +706,10 @@ def generate(
629706 fast_factor : int = 2 ,
630707 fast_sharpen : float = 0.6 ,
631708 rife_weights_dir : str | Path | None = None ,
709+ fast_spatial : bool = False ,
710+ fast_spatial_scale : int = 2 ,
711+ fast_spatial_upsample_mode : str = DEFAULT_PIXEL_UPSAMPLE_MODE ,
712+ fast_spatial_sharpen : float = DEFAULT_FAST_SPATIAL_SHARPEN ,
632713 vsa : bool = False ,
633714 vsa_sparsity : float = 0.9 ,
634715 vsa_tile_size : int = 64 ,
@@ -654,12 +735,23 @@ def generate(
654735 ) if vsa else MiniMaxH3VSAConfig ()
655736 if vsa_config .enabled and not mlx_h3_checkpoint_vsa_capable (self .dit_checkpoint ):
656737 raise dense_only_vsa_error (self .dit_checkpoint )
738+ spatial_plan = plan_fast_spatial (
739+ height ,
740+ width ,
741+ scale = fast_spatial_scale ,
742+ upsample_mode = fast_spatial_upsample_mode ,
743+ sharpen = fast_spatial_sharpen ,
744+ ) if fast_spatial else None
657745 _preflight_media_dependencies (
658746 fast = fast ,
659747 fast_sharpen = fast_sharpen ,
660748 rife_weights_dir = rife_weights_dir ,
749+ fast_spatial = fast_spatial ,
661750 )
662- canvas_height , canvas_width = _model_canvas_size (height , width )
751+ if spatial_plan is not None :
752+ canvas_height , canvas_width = spatial_plan .canvas_height , spatial_plan .canvas_width
753+ else :
754+ canvas_height , canvas_width = _model_canvas_size (height , width )
663755 target_geometry = self .resolve_geometry (canvas_height , canvas_width , num_frames )
664756 fast_plan = plan_fast_temporal (target_geometry ["num_frames" ], fast_factor ) if fast else None
665757 video_num_frames = fast_plan .source_frames if fast_plan is not None else target_geometry ["num_frames" ]
@@ -671,7 +763,7 @@ def generate(
671763 enforce_duration = fast_plan is None ,
672764 )
673765 logger .info (
674- "Geometry: output=%dx%dx%d model=%dx%dx%d audio_frames=%d fast=%s" ,
766+ "Geometry: output=%dx%dx%d model=%dx%dx%d audio_frames=%d fast=%s fast_spatial=%s " ,
675767 width ,
676768 height ,
677769 target_geometry ["num_frames" ],
@@ -680,6 +772,7 @@ def generate(
680772 video_geometry ["num_frames" ],
681773 target_geometry ["num_frames" ],
682774 fast_plan ,
775+ spatial_plan ,
683776 )
684777
685778 _reset_peak_memory ()
@@ -718,7 +811,10 @@ def generate(
718811 num_frames = video_geometry ["num_frames" ],
719812 tiled = tiled_video_decode ,
720813 )
721- frames = _center_crop_frames (frames , height , width )
814+ if spatial_plan is not None :
815+ frames = _center_crop_frames (frames , spatial_plan .stage1_height , spatial_plan .stage1_width )
816+ else :
817+ frames = _center_crop_frames (frames , height , width )
722818 timings ["video_decode_s" ] = time .perf_counter () - started
723819 peaks ["video_decode_gib" ] = _peak_memory_gib ()
724820 _cleanup_mlx ()
@@ -735,7 +831,8 @@ def generate(
735831 target_geometry ["num_frames" ],
736832 model = model ,
737833 )
738- interpolated = _sharpen_frames (interpolated , fast_sharpen )
834+ if spatial_plan is None :
835+ interpolated = _sharpen_frames (interpolated , fast_sharpen )
739836 frames = np .stack (interpolated )
740837 if frames .shape [0 ] != target_geometry ["num_frames" ]:
741838 raise RuntimeError (
@@ -748,6 +845,22 @@ def generate(
748845 del model
749846 _cleanup_mlx ()
750847
848+ if spatial_plan is not None :
849+ started = time .perf_counter ()
850+ # One sharpen pass, at full resolution: RIFE and the resample soften
851+ # for the same reason, so the stronger requested amount is applied
852+ # once instead of stacking two unsharp masks.
853+ sharpen = spatial_plan .sharpen if fast_plan is None else max (spatial_plan .sharpen , fast_sharpen )
854+ frames = np .stack (
855+ upsample_frames (
856+ frames ,
857+ width = spatial_plan .target_width ,
858+ height = spatial_plan .target_height ,
859+ mode = spatial_plan .upsample_mode ,
860+ sharpen = sharpen ,
861+ ))
862+ timings ["spatial_upsample_s" ] = time .perf_counter () - started
863+
751864 _reset_peak_memory ()
752865 started = time .perf_counter ()
753866 waveform = self .decode_audio (audio_rows , num_frames = target_geometry ["num_frames" ])
@@ -759,8 +872,8 @@ def generate(
759872 video_path = self .mux (frames , waveform , output_path )
760873 timings ["mux_s" ] = time .perf_counter () - started
761874 timings ["generate_s" ] = sum (
762- timings .get (key , 0.0 )
763- for key in ( "condition_s" , "denoise_s" , "video_decode_s" , "rife_s ""audio_decode_s" , "mux_s" ))
875+ timings .get (key , 0.0 ) for key in ( "condition_s" , "denoise_s" , "video_decode_s" , "rife_s" ,
876+ "spatial_upsample_s ""audio_decode_s" , "mux_s" ))
764877
765878 result = GenerationResult (
766879 video_path = str (video_path ),
0 commit comments