sg_proc_arguments.py

from __future__ import annotations

from gi.repository import Gimp, GObject

from sg_constants import INSERT_MODES
from sg_i18n import _
from sg_utils import make_choice_from_list


def PLUGIN_FIELDS_CHECKPOINT(
    procedure: Gimp.Procedure, models: list[str], selected_model: str | None, sd_modules: list[str],
) -> None:
    procedure.add_choice_argument(
        "model",
        _("Model"),
        _("Model to use for inference"),
        make_choice_from_list(models),
        selected_model if selected_model in models else models[0],
        GObject.ParamFlags.READWRITE,
    )

    FLUX_ENCODERS_MODES = ["Autoguess", "Never add", "Always add"]
    procedure.add_choice_argument(
        "flux_encoders_mode",
        _("Add VAE and text encoders for Flux"),
        _("""Add VAE and text encoders for Flux-based models:

          Autoguess: Flux VAEs and text encoders will be added automatically based on model name
          Never add: Flux VAEs and text encoders will never be added
          Always add: Flux VAEs and text encoders will always be added"""),
        make_choice_from_list(FLUX_ENCODERS_MODES),
        FLUX_ENCODERS_MODES[0],
        GObject.ParamFlags.READWRITE,
    )

    guessed_modules = [
        module
        for module in sd_modules
        if (
            module.rsplit("/", 1)[-1]
            in (
                "ae.safetensors",
                "clip_l.safetensors",
                "t5xxl_fp8_e4m3fn.safetensors",
            )
            or "flux" in module.lower()
            and ("/vae/" in module.lower() or "/text_encoder/" in module.lower())
        )
    ]

    procedure.add_string_argument(
        "flux_encoders",
        _("Flux VAEs and text encoders"),
        _("Flux VAEs and text encoders to use for Flux-based models"),
        "\n".join(guessed_modules),
        GObject.ParamFlags.READWRITE,
    )


def PLUGIN_FIELDS_COMMON(procedure: Gimp.Procedure, samplers: list[str], selected_sampler: str) -> None:
    procedure.add_string_argument(
        "prompt",
        _("Prompt"),
        _("A text string that describes the image you want to generate"),
        "",
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_string_argument(
        "negative_prompt",
        _("Negative Prompt"),
        _("A text description on what you don't want to see in the image "),
        "",
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_int_argument(
        "seed",
        _("Seed"),
        _("Value that determines the output of a random number generator (0 or -1 for random)"),
        0,
        65535,
        0,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_int_argument(
        "batch_size",
        _("Batch count"),
        _("Specifies the number of results you want to get."),
        1,
        20,
        1,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_int_argument(
        "steps",
        _("Steps"),
        _("""The number of denoising steps used by the AI model when generating an image from the text prompt.
        A higher number results in longer generation time, and does not necessarily guarantee higher image quality."""),
        5,
        150,
        20,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_int_argument(
        "mask_blur",
        _("Mask Blur"),
        _("Feathering of a mask (from edges to inside the mask)"),
        0,
        64,
        4,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_int_argument(
        "width",
        _("Width"),
        _("Width of the image you want to generate"),
        64,
        2048,
        512,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_int_argument(
        "height",
        _("Height"),
        _("Height of the image that you want to generate"),
        64,
        2048,
        512,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_double_argument(
        "cfg_scale",
        _("CFG Scale"),
        _("""Governs how closely it follows the prompt. Lower values produce more creative results.
7 or 7.5 usually works best. You can increase the value if the generated image doesn’t match your prompt.

Low numbers (0-6 ish): You're telling SD that it can ignore your prompt.
Mid-range (6-10 ish): You're telling SD you'd like it to do what you're asking, but you don't mind a bit of variation.
High numbers (10+): We found that high cfg often results in a really bad image."""),
        1,
        30,
        7,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_double_argument(
        "denoising_strength",
        _("Denoising Strength"),
        _("""Strength of image transfomation during inpainting precess.

Determines how little respect the algorithm should have for image's content.
At 0, nothing will change, and at 1 you'll get an unrelated image.

High means more influence during transformation.
A value between 0.5 and 1 is optimal for maintaining image consistency.


0.0: No change
0.35: Slight change
0.75: Quite a change
1.0: Most changes
"""),
        0,
        1,
        0.75,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_choice_argument(
        "sampler_index",
        _("Sampler"),
        _("Sampler to use"),
        make_choice_from_list(samplers),
        selected_sampler,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_boolean_argument(
        "restore_faces",
        _("Restore faces"),
        _("Use GFPGAN / CodeFormer to restore faces"),
        False,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_boolean_argument(
        "tiling",
        _("Tiling"),
        _("Tiling will improve image quality and speed, but may cause artifacts"),
        False,
        GObject.ParamFlags.READWRITE,
    )

    procedure.add_choice_argument(
        "insert_mode",
        _("Insert mode"),
        _("""How to insert generated image:

Scale to fill — generated image heights and widths are stretched to match the size of selection (default)
Insert as is — insert generated image with generation size at selection left-top coords
Aspect fill — longest side of generated image is proportionnaly stretched to match the selection
Aspect fit — shortest side of generated image is proportionnaly stretched to match the selection
Use selection size — try to generate image of selection size (instead of input sizes)
        """),
        make_choice_from_list(INSERT_MODES),
        INSERT_MODES[0],
        GObject.ParamFlags.READWRITE,
    )


def PLUGIN_FIELDS_CONTROLNET_OPTIONS(procedure: Gimp.Procedure) -> None:
    procedure.add_boolean_argument(
        "cn1_enabled",
        _("Enable ControlNet 1"),
        _("Enable ControlNet 1 for generation"),
        False,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_layer_argument(
        "cn1_layer",
        _("ControlNet 1 Layer"),
        _("ControlNet 1 Layer for generation"),
        True,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_boolean_argument(
        "cn2_enabled",
        _("Enable ControlNet 2"),
        _("Enable ControlNet 2 for generation"),
        False,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_layer_argument(
        "cn2_layer",
        _("ControlNet 2 Layer"),
        _("ControlNet 2 Layer for generation"),
        True,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_boolean_argument(
        "cn_skip_annotator_layers",
        _("Skip annotator layers"),
        _("Skip annotator layers (something that I don't know what it does)"),
        True,
        GObject.ParamFlags.READWRITE,
    )


def PLUGIN_FIELDS_RESIZE_MODE(procedure: Gimp.Procedure, resize_modes: list[str]) -> None:
    procedure.add_choice_argument(
        "resize_mode",
        _("Resize Mode"),
        _("""Resize Mode

Just resize: Simply resizes the image to target dimensions.The aspect ratio is not maintained, so the image will stretch or shrink if the aspect ratio of the destination size is different.

Crop and resize: Crops the image to maintain aspect ratio, then resizes. The aspect ratio is maintained and the overhang is removed.

Resize and fill: Resizes and fills empty areas with content. The image is scaled to fit the specified destination size while maintaining the aspect ratio. The excess area will be filled with the color of the edge of the input image.

Just resize (latent upscale): Resizes in the latent space. Same as Just Resize, but uses Latent Upscale for scaling.
        """),  # noqa: E501
        make_choice_from_list(resize_modes),
        resize_modes[0],
        GObject.ParamFlags.READWRITE,
    )


def PLUGIN_FIELDS_INPAINTING(procedure: Gimp.Procedure, inpaint_fill_modes: list[str]) -> None:
    procedure.add_boolean_argument(
        "invert_mask",
        _("Invert Mask"),
        _("Determines whether you want to inpaint the masked object or outside of the masked object"),
        False,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_boolean_argument(
        "inpaint_full_res",
        _("Inpaint Whole Picture"),
        _("""Refers to the inpaint area. It can be either true or false.

        If set to true, inpaint area will be the same resolution as the starting image.
        If set to false, inpaint area will stretch to fit the starting image."""),
        True,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_choice_argument(
        "inpainting_fill",
        _("Mask fill mode"),
        _("""Choose the fill content in mask

fill: Fills the input image with the average color within the masked area. Good for removing objects.
original: The input image is used as is.
latent noise: Fills the masked area with Seed-based noise. It can be used to mask an object and rewrite the background. 1 is recommended for CGF.
latent nothing: Fills the masked area with the average color of the masked area. It can be used to mask an object and rewrite the background. 0.8 or higher is recommended for CGF.
        """),  # noqa: E501
        make_choice_from_list(inpaint_fill_modes),
        inpaint_fill_modes[0],
        GObject.ParamFlags.READWRITE,
    )


def PLUGIN_FIELDS_CONTROLNET(
    procedure: Gimp.Procedure,
    cn_modules: list[str],
    cn_models: list[str],
    cn_resize_modes: list[str],
    control_modes: list[str],
) -> None:
    procedure.add_choice_argument(
        "module",
        _("Module"),
        _("Module to use for controlnet"),
        make_choice_from_list(cn_modules),
        cn_modules[0],
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_choice_argument(
        "model",
        _("Model"),
        _("Model to use for controlnet"),
        make_choice_from_list(cn_models),
        cn_models[0],
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_double_argument(
        "weight",
        _("Weight"),
        _("""Determines how closely the control map should follow the prompt.
By default it is 1 and it will deviate from the control map as it gets closer to 0"""),
        0,
        2,
        1,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_choice_argument(
        "resize_mode",
        _("Resize Mode"),
        _("Resize Mode for ControlNet"),
        make_choice_from_list(cn_resize_modes),
        cn_resize_modes[0],
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_boolean_argument(
        "lowvram",
        _("Low VRAM"),
        _("Only recommended if you have less than 8gb VRAM"),
        False,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_choice_argument(
        "control_mode",
        _("Control Mode"),
        _("Control Mode for ControlNet"),
        make_choice_from_list(control_modes),
        control_modes[0],
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_double_argument(
        "guidance_start",
        _("Guidance Start (T)"),
        _("Guidance Start (T) for ControlNet"),
        0,
        1,
        0,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_double_argument(
        "guidance_end",
        _("Guidance End (T)"),
        _("Guidance End (T) for ControlNet"),
        0,
        1,
        1,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_double_argument(
        "guidance", _("Guidance"), _("Guidance for ControlNet"), 0, 1, 1, GObject.ParamFlags.READWRITE,
    )
    procedure.add_int_argument(
        "processor_res",
        _("Processor Resolution"),
        _("set this to 512. There’s an ongoing bug that results in the cv2 error if it’s not set"),
        64,
        2048,
        512,
        GObject.ParamFlags.READWRITE,
    )
    procedure.add_int_argument(
        "threshold_a", _("Threshold A"), _("Threshold A for ControlNet"), 64, 2048, 64, GObject.ParamFlags.READWRITE,
    )
    procedure.add_int_argument(
        "threshold_b", _("Threshold B"), _("Threshold B for ControlNet"), 64, 2048, 64, GObject.ParamFlags.READWRITE,
    )