90 lines
4.4 KiB
Python
90 lines
4.4 KiB
Python
# pylint: disable=missing-module-docstring,disable=missing-class-docstring,missing-function-docstring,line-too-long
|
|
from comfy_api.latest import io
|
|
|
|
from .helpers import Resolution, HIRES_RATIO
|
|
from .node_utilities import ResolutionParam
|
|
from .flux2 import get_best_valid_resolution as get_flux2_best_valid_resolution
|
|
from .flux import get_best_valid_resolution as get_flux_best_valid_resolution
|
|
from .qwenimage import get_best_valid_resolution as get_qwenimage_best_valid_resolution
|
|
from .sdxl import get_best_valid_resolution as get_sdxl_best_valid_resolution
|
|
|
|
models = ["FLUX.2-dev", "Qwen-Image", "FLUX.1-dev", "SDXL"]
|
|
|
|
|
|
class ResolutionAdvisor(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="IG1ResolutionAdvisor",
|
|
display_name="Resolution Advisor",
|
|
category="IG1 Tools",
|
|
description=f"""From a user input desired resolution, this node will compute and output:
|
|
1. A first generation pass resolution, accounting for model's min len, max size and patch len, with a ratio the closest possible to input resolution.
|
|
2. A boolean indicating if a second pass {HIRES_RATIO}x upscale (HiRes fix recommended) is necessary, aka if the generate resolution is lower than the reference resolution.
|
|
3. A boolean indicating if a third pass (pure upscale) is necessary, aka if the the size post 2nd pass (doubling the resolution) is still under the reference resolution.
|
|
""",
|
|
inputs=[
|
|
ResolutionParam.Input(
|
|
"resolution",
|
|
tooltip="The desired resolution of the image to be generated. Will be adjusted as reference to match patch length if necessary.",
|
|
),
|
|
io.Combo.Input(
|
|
"model",
|
|
options=models,
|
|
default=models[0],
|
|
tooltip="The model you want to compute advises for. This will be used to get the patch length, min lengths and max size.",
|
|
)
|
|
],
|
|
outputs=[
|
|
ResolutionParam.Output(
|
|
"generate",
|
|
display_name="GENERATE",
|
|
tooltip="The first pass generation resolution respecting model's patch len, min lenghts and max size."
|
|
|
|
),
|
|
io.Boolean.Output(
|
|
"hires",
|
|
display_name="NEED_HIRES",
|
|
tooltip=f"Indicate if a second pass, {HIRES_RATIO}x HiRes upscale is needed. True if the generate resolution is lower than the reference resolution."
|
|
),
|
|
io.Boolean.Output(
|
|
"upscale",
|
|
display_name="NEED_UPSCALE",
|
|
tooltip="Indicate if a third pass, regular upscale is needed. True if the HiRes second phase resolution is lower than the reference resolution."
|
|
)
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, resolution, model) -> io.NodeOutput:
|
|
# Compute the flux first pass generation resolution
|
|
# and the adjusted (if necessary) reference resolution.
|
|
if model == "FLUX.2-dev":
|
|
generate_reso = get_flux2_best_valid_resolution(resolution)
|
|
elif model == "Qwen-Image":
|
|
generate_reso = get_qwenimage_best_valid_resolution(resolution)
|
|
elif model == "FLUX.1-dev":
|
|
generate_reso = get_flux_best_valid_resolution(resolution)
|
|
elif model == "SDXL":
|
|
generate_reso = get_sdxl_best_valid_resolution(resolution)
|
|
else:
|
|
ValueError(f"Model f{model} has no internal configuration.")
|
|
# Compute if a HiRes fix x2 second pass is needed to get to the reference resolution
|
|
need_hires = False
|
|
need_upscale = False
|
|
if generate_reso.width < resolution.width or generate_reso.height < resolution.height:
|
|
need_hires = True
|
|
hires = Resolution(
|
|
width=generate_reso.width * HIRES_RATIO,
|
|
height=generate_reso.height * HIRES_RATIO
|
|
)
|
|
# And if a 3rd pass pure upscale is necessary post hires fix
|
|
if hires.width < resolution.width or hires.height < resolution.height:
|
|
need_upscale = True
|
|
# Return to the user everything he needs for next steps
|
|
return io.NodeOutput(
|
|
generate_reso,
|
|
need_hires,
|
|
need_upscale,
|
|
)
|