Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import matplotlib.pyplot as plt | |
| from matplotlib.patches import Patch | |
| import io | |
| import json | |
| from PIL import Image | |
| from typing import List | |
| from scripts.enums import StableDiffusionVersion | |
| from scripts.global_state import get_sd_version | |
| from scripts.ipadapter.weight import calc_weights | |
| INPUT_BLOCK_COLOR = "#61bdee" | |
| MIDDLE_BLOCK_COLOR = "#e2e2e2" | |
| OUTPUT_BLOCK_COLOR = "#dc6e55" | |
| def get_bar_colors( | |
| sd_version: StableDiffusionVersion, input_color, middle_color, output_color | |
| ): | |
| middle_block_idx = 4 if sd_version == StableDiffusionVersion.SDXL else 6 | |
| def get_color(idx): | |
| if idx < middle_block_idx: | |
| return input_color | |
| elif idx == middle_block_idx: | |
| return middle_color | |
| else: | |
| return output_color | |
| return [get_color(i) for i in range(sd_version.transformer_block_num)] | |
| def plot_weights( | |
| numbers: List[float], | |
| colors: List[str], | |
| ): | |
| # Create a bar chart | |
| plt.figure(figsize=(8, 4)) | |
| plt.bar(range(len(numbers)), numbers, color=colors) | |
| plt.xlabel("Transformer Index") | |
| plt.ylabel("Weight") | |
| plt.legend( | |
| handles=[ | |
| Patch(color=color, label=label) | |
| for color, label in ( | |
| (INPUT_BLOCK_COLOR, "Input Block"), | |
| (MIDDLE_BLOCK_COLOR, "Middle Block"), | |
| (OUTPUT_BLOCK_COLOR, "Output Block"), | |
| ) | |
| ], | |
| loc="best", | |
| ) | |
| # Save the plot to a BytesIO buffer | |
| buffer = io.BytesIO() | |
| plt.savefig(buffer, format="png") | |
| plt.close() | |
| buffer.seek(0) | |
| # Convert the buffer to a PIL image and return it | |
| image = Image.open(buffer) | |
| return image | |
| class AdvancedWeightControl: | |
| def __init__(self): | |
| self.group = None | |
| self.weight_type = None | |
| self.weight_plot = None | |
| self.weight_editor = None | |
| self.weight_composition = None | |
| def render(self): | |
| with gr.Group(visible=False) as self.group: | |
| with gr.Row(): | |
| self.weight_type = gr.Dropdown( | |
| choices=[ | |
| "normal", | |
| "ease in", | |
| "ease out", | |
| "ease in-out", | |
| "reverse in-out", | |
| "weak input", | |
| "weak output", | |
| "weak middle", | |
| "strong middle", | |
| "style transfer", | |
| "composition", | |
| "strong style transfer", | |
| "style and composition", | |
| "strong style and composition", | |
| ], | |
| label="Weight Type", | |
| value="normal", | |
| ) | |
| self.weight_composition = gr.Slider( | |
| label="Composition Weight", | |
| minimum=0, | |
| maximum=2.0, | |
| value=1.0, | |
| step=0.01, | |
| visible=False, | |
| ) | |
| self.weight_editor = gr.Textbox(label="Weights", visible=False) | |
| self.weight_plot = gr.Image( | |
| value=None, | |
| label="Weight Plot", | |
| interactive=False, | |
| visible=False, | |
| ) | |
| def register_callbacks( | |
| self, | |
| weight_input: gr.Slider, | |
| advanced_weighting: gr.State, | |
| control_type: gr.Radio, | |
| update_unit_counter: gr.Number, | |
| ): | |
| def advanced_weighting_supported(control_type: str) -> bool: | |
| return control_type in ("IP-Adapter", "Instant-ID") | |
| self.weight_type.change( | |
| fn=lambda weight_type: gr.update( | |
| visible=weight_type | |
| in ("style and composition", "strong style and composition") | |
| ), | |
| inputs=[self.weight_type], | |
| outputs=[self.weight_composition], | |
| ) | |
| def update_weight_textbox( | |
| control_type: str, | |
| weight_type: str, | |
| weight: float, | |
| weight_composition: float, | |
| ): | |
| if not advanced_weighting_supported(control_type): | |
| return gr.update() | |
| sd_version = get_sd_version() | |
| weights = calc_weights(weight_type, weight, sd_version, weight_composition) | |
| return gr.update(value=str([round(w, 2) for w in weights]), visible=True) | |
| trigger_inputs = [self.weight_type, weight_input, self.weight_composition] | |
| for trigger_input in trigger_inputs: | |
| trigger_input.change( | |
| fn=update_weight_textbox, | |
| inputs=[ | |
| control_type, | |
| self.weight_type, | |
| weight_input, | |
| self.weight_composition, | |
| ], | |
| outputs=[self.weight_editor], | |
| ) | |
| def update_plot(weights_string: str): | |
| try: | |
| weights = json.loads(weights_string) | |
| assert isinstance(weights, list) | |
| except Exception: | |
| return gr.update(visible=False) | |
| sd_version = get_sd_version() | |
| weight_plot = plot_weights( | |
| weights, | |
| get_bar_colors( | |
| sd_version, | |
| input_color=INPUT_BLOCK_COLOR, | |
| middle_color=MIDDLE_BLOCK_COLOR, | |
| output_color=OUTPUT_BLOCK_COLOR, | |
| ), | |
| ) | |
| return gr.update(value=weight_plot, visible=True) | |
| def update_advanced_weighting(weights_string: str): | |
| try: | |
| weights = json.loads(weights_string) | |
| assert isinstance(weights, list) | |
| except Exception: | |
| return None | |
| return weights | |
| self.weight_editor.change( | |
| fn=update_plot, | |
| inputs=[self.weight_editor], | |
| outputs=[self.weight_plot], | |
| ) | |
| self.weight_editor.change( | |
| fn=update_advanced_weighting, | |
| inputs=[self.weight_editor], | |
| outputs=[advanced_weighting], | |
| ).then( | |
| fn=lambda x: gr.update(value=x + 1), | |
| inputs=[update_unit_counter], | |
| outputs=[update_unit_counter], | |
| ) # Necessary to flush gr.State change to unit state. | |
| # TODO: Expose advanced weighting control for other control types. | |
| def control_type_change(control_type: str, old_weights): | |
| supported = advanced_weighting_supported(control_type) | |
| if supported: | |
| return ( | |
| gr.update(visible=supported), | |
| old_weights, | |
| gr.update(), | |
| gr.update(), | |
| ) | |
| else: | |
| return ( | |
| gr.update(visible=supported), | |
| None, | |
| gr.update(visible=False), | |
| gr.update(visible=False), | |
| ) | |
| control_type.change( | |
| fn=control_type_change, | |
| inputs=[control_type, advanced_weighting], | |
| outputs=[ | |
| self.group, | |
| advanced_weighting, | |
| self.weight_editor, | |
| self.weight_plot, | |
| ], | |
| ) | |