From 773f980a5b2eafb1ded0d8bcaba0051c5e7a892f Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Wed, 4 Feb 2026 14:23:23 +1100 Subject: [PATCH 1/2] replace textbox input for refiner plan with dynamic sliders ui --- requirements.txt | 1 + shared/gradio/ui_styles.css | 11 ++++++ shared/utils/self_refiner.py | 63 ++++++++++------------------- wgp.py | 77 ++++++++++++++++++++++++++++++++---- 4 files changed, 102 insertions(+), 50 deletions(-) diff --git a/requirements.txt b/requirements.txt index f37df8c02..c2c8d2924 100644 --- a/requirements.txt +++ b/requirements.txt @@ -35,6 +35,7 @@ s3tokenizer conformer==0.3.2 spacy_pkuseg spacy==3.8.4 +gradio_rangeslider # Vision & segmentation opencv-python>=4.12.0.88 diff --git a/shared/gradio/ui_styles.css b/shared/gradio/ui_styles.css index f55edf64e..ec3d90cdc 100644 --- a/shared/gradio/ui_styles.css +++ b/shared/gradio/ui_styles.css @@ -364,3 +364,14 @@ user-select: text; overflow: visible !important; } .tabitem {padding-top:0px} +.rule-row { margin-bottom: 8px; align-items: center !important; display: flex; gap: 8px; } +.rule-card { background-color: var(--background-fill-secondary); padding: 8px 12px; border-radius: 6px; border: 1px solid var(--border-color-primary); flex-grow: 1; margin-bottom: 0 !important; } +.rule-card p { margin-bottom: 0 !important; } +.delete-btn { + min-width: 42px !important; + max-width: 42px !important; + height: 42px !important; + padding: 0 !important; + align-self: center; +} +#refiner-input-row { align-items: center; } \ No newline at end of file diff --git a/shared/utils/self_refiner.py b/shared/utils/self_refiner.py index 8ef60abf8..3f6b1ef4d 100644 --- a/shared/utils/self_refiner.py +++ b/shared/utils/self_refiner.py @@ -9,39 +9,22 @@ def is_int_string(s: str) -> bool: except ValueError: return False -def normalize_self_refiner_plan(plan_str): - entries = [] - for chunk in plan_str.split(","): - chunk = chunk.strip() - if not chunk: - continue - if ":" not in chunk: - return "", "Self-refine plan entries must be in 'start-end:steps' format." - range_part, steps_part = chunk.split(":", 1) - range_part = range_part.strip() - steps_part = steps_part.strip() - if not steps_part: - return "", "Self-refine plan entries must include a step count." - if "-" in range_part: - start_s, end_s = range_part.split("-", 1) - else: - start_s = end_s = range_part - start_s = start_s.strip() - if not is_int_string(start_s): - return "", "Self-refine plan start position must be an integer." - end_s = end_s.strip() - if not is_int_string(end_s): - return "", "Self-refine plan end position must be an integer." - if not is_int_string(steps_part): - return "", "Self-refine plan steps part must be an integer." - - entries.append({ - "start": int(start_s), - "end": int(end_s), - "steps": int(steps_part), - }) - plan = entries - return plan, "" +def normalize_self_refiner_plan(plan_input): + default_plan = [ + {"start": 2, "end": 5, "steps": 3}, + {"start": 6, "end": 13, "steps": 1} + ] + + if not plan_input: + return default_plan, "" + + if isinstance(plan_input, list): + for entry in plan_input: + if not all(k in entry for k in ("start", "end", "steps")): + return [], "Self-refine plan entries must contain start, end, and steps." + return plan_input, "" + + return [], "Invalid plan format (expected list)." class PnPHandler: def __init__(self, stochastic_plan, ths_uncertainty=0.0, p_norm=1, certain_percentage=0.999): @@ -343,13 +326,9 @@ def restore_func(saved_state): return latents, sample_scheduler def create_self_refiner_handler(pnp_plan, pnp_f_uncertainty, pnp_p_norm, pnp_certain_percentage): - if len(pnp_plan): - stochastic_plan, error = normalize_self_refiner_plan(pnp_plan) - else: - # Default plan from paper/code - stochastic_plan = [ - {"start": 1, "end": 5, "steps": 3}, - {"start": 6, "end": 13, "steps": 1}, - ] + stochastic_plan, error = normalize_self_refiner_plan(pnp_plan) + if error: + print(f"Warning: Self-refiner plan error: {error}. Using defaults.") + stochastic_plan, _ = normalize_self_refiner_plan(None) - return PnPHandler(stochastic_plan, ths_uncertainty=pnp_f_uncertainty, p_norm=pnp_p_norm, certain_percentage=pnp_certain_percentage) + return PnPHandler(stochastic_plan, ths_uncertainty=pnp_f_uncertainty, p_norm=pnp_p_norm, certain_percentage=pnp_certain_percentage) \ No newline at end of file diff --git a/wgp.py b/wgp.py index 135b42c7c..9e36c58e1 100644 --- a/wgp.py +++ b/wgp.py @@ -65,6 +65,8 @@ import glob import cv2 import html +import uuid +from gradio_rangeslider import RangeSlider from transformers.utils import logging logging.set_verbosity_error from tqdm import tqdm @@ -610,12 +612,11 @@ def ret(): model_mode = None if server_config.get("fit_canvas", 0) == 2 and outpainting_dims is not None and any_letters(video_prompt_type, "VKF"): gr.Info("Output Resolution Cropping will be not used for this Generation as it is not compatible with Video Outpainting") - if self_refiner_setting != 0 and len(self_refiner_plan): - from shared.utils.self_refiner import normalize_self_refiner_plan - _, error = normalize_self_refiner_plan(self_refiner_plan) - if len(error): - gr.Info(error) - return ret() + if self_refiner_setting != 0: + if isinstance(self_refiner_plan, list): + pass + elif self_refiner_plan: + self_refiner_plan = [] if not model_def.get("motion_amplitude", False): motion_amplitude = 1. if "vae" in spatial_upsampling: @@ -9150,6 +9151,36 @@ def download_lora(state, lora_url, progress=gr.Progress(track_tqdm=True),): def set_gallery_tab(state, evt:gr.SelectData): return evt.index, "video" if evt.index == 0 else "audio" +def ensure_refiner_list(plan_data): + if not isinstance(plan_data, list): + return [] + for rule in plan_data: + if "id" not in rule: + rule["id"] = str(uuid.uuid4()) + return plan_data + +def add_refiner_rule(current_rules, range_val, steps_val): + new_start, new_end = int(range_val[0]), int(range_val[1]) + + if new_start >= new_end: + raise gr.Error(f"Start step ({new_start}) must be smaller than End step ({new_end}).") + + for rule in current_rules: + if new_start <= rule['end'] and new_end >= rule['start']: + raise gr.Error(f"Overlap detected! Steps {new_start}-{new_end} conflict with existing rule {rule['start']}-{rule['end']}.") + + new_rule = { + "id": str(uuid.uuid4()), + "start": new_start, + "end": new_end, + "steps": int(steps_val) + } + updated_list = current_rules + [new_rule] + return sorted(updated_list, key=lambda x: x['start']) + +def remove_refiner_rule(current_rules, rule_id): + return [r for r in current_rules if r["id"] != rule_id] + def generate_video_tab(update_form = False, state_dict = None, ui_defaults = None, model_family = None, model_base_type_choice = None, model_choice = None, header = None, main = None, main_tabs= None, tab_id='generate', edit_tab=None, default_state=None): global inputs_names #, advanced plugin_data = gr.State({}) @@ -10137,8 +10168,38 @@ def gen_upsampling_dropdowns(temporal_upsampling, spatial_upsampling , film_grai with gr.Column(visible = model_def.get("self_refiner", False)) as self_refiner_col: gr.Markdown("Self-Refining Video Sampling (PnP) - should improve quality of Motion") - self_refiner_setting = gr.Dropdown( choices=[("Disabled", 0),("Enabled with P1-Norm", 1), ("Enabled with P2-Norm", 2), ], value=ui_get("self_refiner_setting", 0), scale = 1, label="Self Refiner", ) - self_refiner_plan = gr.Textbox( value=ui_get("self_refiner_plan", ""), label="P&P Plan (start-end:steps, comma-separated)", lines=1, placeholder="2-5:3,6-13:1" ) + self_refiner_setting = gr.Dropdown(choices=[("Disabled", 0),("Enabled with P1-Norm", 1), ("Enabled with P2-Norm", 2)], value=ui_get("self_refiner_setting", 0), scale=1, label="Self Refiner") + + refiner_val = ensure_refiner_list(ui_get("self_refiner_plan", [])) + self_refiner_plan = refiner_val if update_form else gr.State(value=refiner_val) + + with gr.Group(visible=(update_form and ui_get("self_refiner_setting", 0) > 0)) as self_refiner_rules_ui: + gr.Markdown("### Refiner Plan") + + with gr.Row(elem_id="refiner-input-row"): + refiner_range = RangeSlider(minimum=0, maximum=100, value=(0, 10), step=1, label="Step Range", info="Start - End", scale=3) + refiner_mult = gr.Slider(label="Iterations", value=3, minimum=1, maximum=5, step=1, scale=2) + refiner_add_btn = gr.Button("➕ Add", variant="primary", scale=0, min_width=100) + + if not update_form: + refiner_add_btn.click(fn=add_refiner_rule, inputs=[self_refiner_plan, refiner_range, refiner_mult], outputs=[self_refiner_plan]) + self_refiner_setting.change(fn=lambda s: gr.update(visible=s > 0), inputs=[self_refiner_setting], outputs=[self_refiner_rules_ui]) + + @gr.render(inputs=self_refiner_plan) + def render_refiner_plans(plans): + if not plans: + gr.Markdown("No plans defined. Using defaults: Steps 2-5 (3x), Steps 6-13 (1x).") + return + for plan in plans: + with gr.Row(elem_classes="rule-row"): + text_display = f"Steps **{plan['start']} - {plan['end']}** : **{plan['steps']}x** iterations" + gr.Markdown(text_display, elem_classes="rule-card") + gr.Button("✖", variant="stop", scale=0, elem_classes="delete-btn").click( + fn=remove_refiner_rule, + inputs=[self_refiner_plan, gr.State(plan["id"])], + outputs=[self_refiner_plan] + ) + # self_refiner_f_uncertainty = gr.Slider(0.0, 1.0, value=ui_get("self_refiner_f_uncertainty", 0.0), step=0.01, label="Uncertainty Threshold", show_reset_button= False) # self_refiner_certain_percentage = gr.Slider(0.0, 1.0, value=ui_get("self_refiner_certain_percentage", 0.999), step=0.001, label="Certainty Percentage Skip", show_reset_button= False) From d8f66126992d59e0f6d6c6fca82db66d67453633 Mon Sep 17 00:00:00 2001 From: Chris Malone Date: Wed, 11 Feb 2026 04:13:42 +1100 Subject: [PATCH 2/2] moved ui functions to self_refiner.py --- shared/utils/self_refiner.py | 33 +++++++++++++++++++++++++++++++++ wgp.py | 33 +-------------------------------- 2 files changed, 34 insertions(+), 32 deletions(-) diff --git a/shared/utils/self_refiner.py b/shared/utils/self_refiner.py index e595b092d..9d64705e0 100644 --- a/shared/utils/self_refiner.py +++ b/shared/utils/self_refiner.py @@ -1,5 +1,6 @@ import torch import copy +import uuid from diffusers.utils.torch_utils import randn_tensor def is_int_string(s: str) -> bool: @@ -21,6 +22,38 @@ def normalize_self_refiner_plan(plan_input, max_plans: int = 1): return [plan_input], "" +def ensure_refiner_list(plan_data): + if not isinstance(plan_data, list): + return [] + for rule in plan_data: + if "id" not in rule: + rule["id"] = str(uuid.uuid4()) + return plan_data + +def add_refiner_rule(current_rules, range_val, steps_val): + new_start, new_end = int(range_val[0]), int(range_val[1]) + + if new_start >= new_end: + from gradio import Error + raise Error(f"Start step ({new_start}) must be smaller than End step ({new_end}).") + + for rule in current_rules: + if new_start <= rule['end'] and new_end >= rule['start']: + from gradio import Error + raise Error(f"Overlap detected! Steps {new_start}-{new_end} conflict with existing rule {rule['start']}-{rule['end']}.") + + new_rule = { + "id": str(uuid.uuid4()), + "start": new_start, + "end": new_end, + "steps": int(steps_val) + } + updated_list = current_rules + [new_rule] + return sorted(updated_list, key=lambda x: x['start']) + +def remove_refiner_rule(current_rules, rule_id): + return [r for r in current_rules if r["id"] != rule_id] + class PnPHandler: def __init__(self, stochastic_plan, ths_uncertainty=0.0, p_norm=1, certain_percentage=0.999, channel_dim: int = 1): self.stochastic_step_map = self._build_stochastic_step_map(stochastic_plan) diff --git a/wgp.py b/wgp.py index e4e4c64e9..96c92d8c5 100644 --- a/wgp.py +++ b/wgp.py @@ -48,6 +48,7 @@ from huggingface_hub import hf_hub_download, snapshot_download from shared.utils import files_locator as fl from shared.gradio.audio_gallery import AudioGallery +from shared.utils.self_refiner import normalize_self_refiner_plan, ensure_refiner_list, add_refiner_rule, remove_refiner_rule import torch import gc import traceback @@ -65,7 +66,6 @@ import glob import cv2 import html -import uuid from gradio_rangeslider import RangeSlider from transformers.utils import logging logging.set_verbosity_error @@ -614,7 +614,6 @@ def ret(): gr.Info("Output Resolution Cropping will be not used for this Generation as it is not compatible with Video Outpainting") if self_refiner_setting != 0: if isinstance(self_refiner_plan, list): - from shared.utils.self_refiner import normalize_self_refiner_plan max_plans = model_def.get("self_refiner_max_plans", 1) _, error = normalize_self_refiner_plan(self_refiner_plan, max_plans=max_plans) if len(error): @@ -9167,36 +9166,6 @@ def download_lora(state, lora_url, progress=gr.Progress(track_tqdm=True),): def set_gallery_tab(state, evt:gr.SelectData): return evt.index, "video" if evt.index == 0 else "audio" -def ensure_refiner_list(plan_data): - if not isinstance(plan_data, list): - return [] - for rule in plan_data: - if "id" not in rule: - rule["id"] = str(uuid.uuid4()) - return plan_data - -def add_refiner_rule(current_rules, range_val, steps_val): - new_start, new_end = int(range_val[0]), int(range_val[1]) - - if new_start >= new_end: - raise gr.Error(f"Start step ({new_start}) must be smaller than End step ({new_end}).") - - for rule in current_rules: - if new_start <= rule['end'] and new_end >= rule['start']: - raise gr.Error(f"Overlap detected! Steps {new_start}-{new_end} conflict with existing rule {rule['start']}-{rule['end']}.") - - new_rule = { - "id": str(uuid.uuid4()), - "start": new_start, - "end": new_end, - "steps": int(steps_val) - } - updated_list = current_rules + [new_rule] - return sorted(updated_list, key=lambda x: x['start']) - -def remove_refiner_rule(current_rules, rule_id): - return [r for r in current_rules if r["id"] != rule_id] - def generate_video_tab(update_form = False, state_dict = None, ui_defaults = None, model_family = None, model_base_type_choice = None, model_choice = None, header = None, main = None, main_tabs= None, tab_id='generate', edit_tab=None, default_state=None): global inputs_names #, advanced plugin_data = gr.State({})