divisor.keybinds

  1# SPDX-License-Identifier: MPL-2.0 AND LicenseRef-Commons-Clause-License-Condition-1.0
  2# <!-- // /*  d a r k s h a p e s */ -->
  3
  4"""
  5Choice registry decorator.
  6
  7Usage:
  8```
  9@choice("g", "Guidance")
 10def change_guidance(controller: ManualTimestepController,
 11state: MenuState,
 12clear_prediction_cache: Callable[[], None],) -> MenuState:
 13
 14    ...
 15
 16    return state
 17```
 18
 19Returns:
 20```
 21{"fn": callable, "desc": description}
 22```
 23
 24"""
 25
 26from collections import OrderedDict
 27from typing import Any, Callable, Dict
 28
 29from nnll.console import nfo
 30from nnll.helpers import generate_valid_resolutions
 31
 32from divisor.cli_input import (
 33    get_float_input,
 34    get_int_input,
 35    handle_float_setting,
 36    handle_toggle,
 37)
 38from divisor.controller import ManualTimestepController, update_state_and_cache
 39from divisor.interaction_context import InteractionContext
 40from divisor.noise import prepare_4d_noise_for_3d_model
 41from divisor.state import MenuState
 42
 43
 44def choice(key: str, description: str) -> Callable[[Callable], Callable]:
 45    """Decorator that registers a function as a menu choice.\n
 46    :param key: Single-character option the user will type (e.g. ``"g"``).Use ``""`` for the *Enter* (advance) option.
 47    :param description: Short human-readable description that will be shown in the prompt and self-document the registry"""
 48
 49    def wrapper(fn: Callable) -> Callable:
 50        normalized_key = key.lower()
 51        _CHOICE_REGISTRY[normalized_key] = {"fn": fn, "desc": description}  # case‑insensitive
 52        return fn
 53
 54    return wrapper
 55
 56
 57_CHOICE_REGISTRY: "OrderedDict[str, Dict[str, Any]]" = OrderedDict()
 58
 59
 60@choice("", "Advance (Enter)")
 61def _advance(
 62    controller: ManualTimestepController,
 63    state: MenuState,
 64    interaction_context: InteractionContext,
 65) -> MenuState:
 66    """Advance one step (default action when Enter is pressed)."""
 67    nfo("Advancing...")
 68    interaction_context.clear_prediction_cache()
 69    controller.step()
 70    return controller.current_state
 71
 72
 73@choice("g", "Guidance")
 74def change_guidance(
 75    controller: ManualTimestepController,
 76    state: MenuState,
 77    interaction_context: InteractionContext,
 78) -> MenuState:
 79    """Handle guidance value change.\n"""
 80    new_guidance = get_float_input(
 81        f"Enter new guidance value (current: {state.guidance:.2f}): ",
 82        state.guidance,
 83        allow_empty=False,
 84    )
 85    if new_guidance is not None:
 86        return update_state_and_cache(
 87            controller,
 88            controller.set_guidance,
 89            new_guidance,
 90            interaction_context,
 91            f"Guidance set to {new_guidance}",
 92        )
 93    nfo("Invalid guidance value, keeping current value")
 94    return state
 95
 96
 97@choice("l", "Layer Dropout")
 98def change_layer_dropout(
 99    controller: ManualTimestepController,
100    state: MenuState,
101    interaction_context: InteractionContext,
102) -> MenuState:
103    """Handle layer dropout change.\n
104    :param controller: ManualTimestepController instance
105    :param state: Current input state
106    :param interaction_context: InteractionContext arguments
107    :returns: Updated input state"""
108    try:
109        dropout_input = input("Enter layer indices to drop (comma-separated, or 'none' to clear): ").strip()
110        if dropout_input.lower() == "none" or dropout_input == "":
111            layer_indices = None
112        else:
113            layer_indices = [int(x.strip()) for x in dropout_input.split(",")]
114
115        controller.set_layer_dropout(layer_indices)
116        interaction_context.clear_prediction_cache()
117        state = controller.current_state
118        nfo(f"Layer dropout set to: {layer_indices}")
119        return state
120    except ValueError:
121        nfo("Invalid layer indices, keeping current value")
122    return controller.current_state
123
124
125@choice("r", "Resolution")
126def change_resolution(
127    controller: ManualTimestepController,
128    state: MenuState,
129    interaction_context: InteractionContext,
130) -> MenuState:
131    """Handle resolution change.\n
132    :param controller: ManualTimestepController instance
133    :param state: Current input state
134    :param clear_prediction_cache: Function to clear prediction cache
135    :returns: Updated input state"""
136
137    try:
138        if state.width is None or state.height is None:
139            nfo("Cannot generate resolutions: width or height not set")
140        else:
141            valid_resolutions = generate_valid_resolutions(state.width, state.height)
142            nfo("Valid resolutions (same patch count):")
143            for i, (w, h) in enumerate(valid_resolutions):
144                current_marker = ""
145                if state.width == w and state.height == h:
146                    current_marker = " (current)"
147                nfo(f"  {i}: {w}x{h}{current_marker}")
148            resolution_input = input(f"\nEnter resolution index (0-{len(valid_resolutions) - 1}) or 'custom' for custom: ").strip()
149            if resolution_input.lower() == "custom":
150                width_input = input("Enter width: ").strip()
151                height_input = input("Enter height: ").strip()
152                new_width = int(width_input)
153                new_height = int(height_input)
154            else:
155                resolution_idx = int(resolution_input)
156                if 0 <= resolution_idx < len(valid_resolutions):
157                    new_width, new_height = valid_resolutions[resolution_idx]
158                else:
159                    nfo("Invalid resolution index, keeping current value")
160                    return state
161            if new_width is not None and new_height is not None:
162                controller.set_resolution(new_width, new_height)
163                interaction_context.clear_prediction_cache()
164                state = controller.current_state
165                nfo(f"Resolution set to: {new_width}x{new_height}")
166                return state
167    except (ValueError, IndexError):
168        nfo("Invalid resolution input, keeping current value")
169    return state
170
171
172@choice("s", "Seed")
173def change_seed(
174    controller: ManualTimestepController,
175    state: MenuState,
176    interaction_context: InteractionContext,
177) -> MenuState:
178    """Handle seed change.\n
179    :param controller: ManualTimestepController instance
180    :param state: Current menu state
181    :param rng: Random number generator instance
182    :param clear_prediction_cache: Function to clear prediction cache
183    :param t5: Optional T5 embedder instance (required for Flux1/XFlux1 to prepare sample)
184    :param clip: Optional CLIP embedder instance (required for Flux1/XFlux1 to prepare sample)
185    :returns: Updated menu state"""
186
187    current_seed = state.seed if state.seed is not None else 0
188    if new_seed := get_int_input(
189        f"Enter new seed number (current: {current_seed}, or press Enter for random): ",
190        current_seed,
191        generate_random=lambda: interaction_context.rng.next_seed(),
192    ):
193        controller.set_seed(new_seed)
194        new_sample = prepare_4d_noise_for_3d_model(
195            height=state.height,  # type: ignore
196            width=state.width,  # type: ignore
197            seed=new_seed,
198            t5=interaction_context.t5,
199            clip=interaction_context.clip,
200            prompt=state.prompt,
201        )
202
203        controller.current_sample = new_sample
204        interaction_context.clear_prediction_cache()
205        updated_state = controller.current_state
206        nfo(f"Seed set to: {new_seed}")
207        return updated_state
208    nfo("Invalid seed value, keeping current seed")
209    return state
210
211
212@choice("b", "Buffer Mask")
213def toggle_buffer_mask(
214    controller: ManualTimestepController,
215    state: MenuState,
216    interaction_context: InteractionContext,
217) -> MenuState:
218    """Handle buffer mask toggle.\n
219    :param controller: ManualTimestepController instance
220    :param state: Current input state
221    :returns: Updated input state"""
222
223    return handle_toggle(
224        controller,
225        state,
226        state.use_previous_as_mask,
227        controller.set_use_previous_as_mask,
228        enabled_msg="Previous step tensor mask: ENABLED",
229        disabled_msg="Previous step tensor mask: DISABLED",
230    )
231
232
233@choice("a", "Autoencoder Offset")
234def change_vae_offset(
235    controller: ManualTimestepController,
236    state: MenuState,
237    interaction_context: InteractionContext,
238) -> MenuState:
239    """Handle VAE shift/scale offset change.\n
240    :param controller: ManualTimestepController instance
241    :param state: Current input state
242    :param ae: Optional AutoEncoder instance
243    :param clear_prediction_cache: Function to clear prediction cache
244    :returns: Updated input state"""
245
246    if interaction_context.ae is None:
247        nfo("AutoEncoder not available, cannot set VAE offset")
248        return state
249
250    vae_input = input("\nChoose [S]hift or [C]scale: ").strip().lower()
251    if vae_input == "c":
252        return handle_float_setting(
253            controller,
254            state,
255            f"Enter VAE scale offset (current: {state.vae_scale_offset:.4f}, or press Enter to reset to 0.0): ",
256            state.vae_scale_offset,
257            controller.set_vae_scale_offset,
258            interaction_context.clear_prediction_cache,
259            default_value=0.0,
260            value_name="VAE scale offset",
261        )
262    elif vae_input == "s":
263        return handle_float_setting(
264            controller,
265            state,
266            f"Enter VAE shift offset (current: {state.vae_shift_offset:.4f}, or press Enter to reset to 0.0): ",
267            state.vae_shift_offset,
268            controller.set_vae_shift_offset,
269            interaction_context.clear_prediction_cache,
270            default_value=0.0,
271            value_name="VAE shift offset",
272        )
273    return state
274
275
276@choice("d", "Deterministic")
277def toggle_deterministic(
278    controller: ManualTimestepController,
279    state: MenuState,
280    interaction_context: InteractionContext,
281) -> MenuState:
282    """Handle deterministic mode toggle.\n
283    :param controller: ManualTimestepController instance
284    :param state: Current input state
285    :param clear_prediction_cache: Function to clear prediction cache
286    :returns: Updated input state"""
287
288    deterministic = not state.deterministic
289    interaction_context.rng.random_mode(reproducible=deterministic)
290    interaction_context.variation_rng.random_mode(reproducible=deterministic)
291
292    return handle_toggle(
293        controller,
294        state,
295        state.deterministic,
296        controller.set_deterministic,
297        interaction_context.clear_prediction_cache,
298        enabled_msg="Deterministic mode: ENABLED",
299        disabled_msg="Deterministic mode: DISABLED",
300    )
301
302
303@choice("p", "Prompt")
304def change_prompt(
305    controller: ManualTimestepController,
306    state: MenuState,
307    interaction_context: InteractionContext,
308) -> MenuState:
309    """Handle prompt change.\n
310    :param controller: ManualTimestepController instance
311    :param state: Current input state
312    :param clear_prediction_cache: Function to clear prediction cache
313    :param recompute_text_embeddings: Optional function to recompute text embeddings
314    :returns: Updated input state"""
315
316    current_prompt = state.prompt if state.prompt is not None else ""
317    new_prompt = input(f"Enter new prompt (current: {current_prompt}): ").strip()
318
319    if new_prompt:
320        controller.set_prompt(new_prompt)
321        if interaction_context.recompute_text_embeddings is not None:
322            interaction_context.recompute_text_embeddings(new_prompt)
323        else:
324            interaction_context.clear_prediction_cache()
325        state = controller.current_state
326        nfo(f"Prompt set to: {new_prompt}")
327    else:
328        nfo("Prompt unchanged")
329
330    return state
331
332
333@choice("e", "Edit Mode")
334def edit_mode(
335    controller: ManualTimestepController,
336    state: MenuState,
337    interaction_context: InteractionContext,
338) -> None:
339    """Handle edit mode (debugger breakpoint).\n
340    :param clear_prediction_cache: Function to clear prediction cache"""
341    nfo("Entering edit mode (use c/cont to exit)...")
342    breakpoint()
343    interaction_context.clear_prediction_cache()
344
345
346@choice("j", "Jump to Step")
347def jump_to_step(
348    controller: ManualTimestepController,
349    state: MenuState,
350    interaction_context: InteractionContext,
351) -> MenuState:
352    """Run the controller forward until a user‑specified step index.\n
353    :param controller: ManualTimestepController instance
354    :param state: Current input state (unused – we return the controller’s final state)
355    :param clear_prediction_cache: Routine to nvalidate cache prediction.
356    :returns: Input state post jump (or the current state if the target is out of range)."""
357
358    target_str = input(f"Jump to step (0‑{controller.current_index} …{len(controller.timesteps) - 1}): ").strip()
359    if not target_str.isdigit():
360        nfo("Invalid step number – aborting jump.")
361        return controller.current_state
362
363    target = int(target_str)
364
365    if target <= controller.current_index:
366        nfo("Target step is before or equal to the current step – nothing to do.")
367        return controller.current_state
368    if target >= len(controller.timesteps):
369        nfo("Target step exceeds the schedule – jumping to the end.")
370        target = len(controller.timesteps) - 1
371
372    interaction_context.clear_prediction_cache()
373    while controller.current_index < target and not controller.is_complete:
374        controller.step()
375
376    nfo(f"Jumped to step {controller.current_index}/{len(controller.timesteps) - 1}")
377    return controller.current_state
378
379
380@choice("v", "Variation")
381def change_variation(
382    controller: ManualTimestepController,
383    state: MenuState,
384    interaction_context: InteractionContext,
385) -> MenuState:
386    """Handle variation seed/strength change.\n
387    :param controller: ManualTimestepController instance
388    :param state: Current input state
389    :param variation_rng: Variation random number generator instance
390    :param clear_prediction_cache: Function to clear prediction cache
391    :returns: Updated input state"""
392    try:
393        var_input = input(
394            f"Variation (current integer seed: {state.variation_seed}, float strength: {state.variation_strength:.3f}. type a number, leave empty for random seed, or use 0.0 to disable): "
395        ).strip()
396
397        if not var_input or "." not in var_input:
398            # Try to parse as integer (seed)
399            try:
400                if var_input != "":
401                    variation_seed = interaction_context.variation_rng.next_seed(int(var_input))
402                else:
403                    variation_seed = interaction_context.variation_rng.next_seed()
404                controller.set_variation_seed(variation_seed)
405                interaction_context.clear_prediction_cache()
406                state = controller.current_state
407                nfo(f"Variation seed set to: {state.variation_seed}")
408            except ValueError:
409                nfo("Invalid integer seed value, keeping current value")
410        else:
411            # Try to parse as float (strength)
412            try:
413                strength_value = float(var_input)
414                if strength_value < 0.0 or strength_value > 1.0:
415                    state = controller.current_state
416                    nfo("Variation strength must be between 0.0 and 1.0, keeping current value")
417                else:
418                    controller.set_variation_strength(strength_value)
419                    interaction_context.clear_prediction_cache()
420                    state = controller.current_state
421                    nfo(f"Variation strength set to: {strength_value:.3f}")
422            except ValueError:
423                nfo("Invalid float strength value, keeping current value")
424    except (ValueError, KeyboardInterrupt):
425        nfo("Invalid variation value, keeping current value")
426    return state
427
428
429@choice("w", "Rewind")
430def rewind(
431    controller: ManualTimestepController,
432    state: MenuState,
433    interaction_context: InteractionContext,
434) -> MenuState:
435    """Rewind the controller by a specified number of steps.\n
436    :param controller: ManualTimestepController instance
437    :param state: Current input state
438    :param clear_prediction_cache: Function to clear prediction cache
439    :returns: Updated input state"""
440    import random as prng
441
442    num_steps = get_int_input(
443        f"Enter number of steps to rewind (current: {controller.rewind_steps}): ",
444        controller.rewind_steps,
445        generate_random=lambda: prng.randint(0, controller.timesteps.index(state.timestep_index)),
446    )
447    if num_steps is not None:
448        controller.rewind(num_steps)
449        interaction_context.clear_prediction_cache()
450        state = controller.current_state
451        nfo(f"Rewind step {num_steps} to {controller.current_index}")
452        return state
453    nfo("Invalid number of steps, keeping current value")
454    return state
455
456
457# Optionally recompute the *remaining* schedule with a different compression
458# remaining = controller.timesteps[controller.current_index + 1 :]
459# new_tail = time_shift(
460#    schedule_mu=controller.mu,
461#    schedule_sigma=controller.sigma,
462#    original_timestep_tensor=torch.tensor(remaining),
463#    desired_step_count=len(remaining),
464#    compression_factor=0.9,               # e.g. stretch the tail a bit
465# )
466# controller.timesteps[controller.current_index + 1 :] = new_tail.tolist()