Skip to content

Commit

Permalink
add support FluxFill inpainting pipeline
Browse files Browse the repository at this point in the history
  • Loading branch information
eaidova committed Dec 26, 2024
1 parent 014a840 commit 58ada07
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 16 deletions.
39 changes: 26 additions & 13 deletions optimum/exporters/openvino/convert.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@
_torch_version,
_transformers_version,
compare_versions,
is_diffusers_version,
is_openvino_tokenizers_version,
is_tokenizers_version,
is_transformers_version,
Expand Down Expand Up @@ -988,24 +989,36 @@ def _get_submodels_and_export_configs(
def get_diffusion_models_for_export_ext(
pipeline: "DiffusionPipeline", int_dtype: str = "int64", float_dtype: str = "fp32", exporter: str = "openvino"
):
try:
from diffusers import (
StableDiffusion3Img2ImgPipeline,
StableDiffusion3InpaintPipeline,
StableDiffusion3Pipeline,
)
if is_diffusers_version(">=", "0.29.0"):
from diffusers import StableDiffusion3Img2ImgPipeline, StableDiffusion3Pipeline

is_sd3 = isinstance(
pipeline, (StableDiffusion3Pipeline, StableDiffusion3InpaintPipeline, StableDiffusion3Img2ImgPipeline)
)
except ImportError:
sd3_pipes = [StableDiffusion3Pipeline, StableDiffusion3Img2ImgPipeline]
if is_diffusers_version(">=", "0.30.0"):
from diffusers import StableDiffusion3InpaintPipeline

sd3_pipes.append(StableDiffusion3InpaintPipeline)

is_sd3 = isinstance(pipeline, tuple(sd3_pipes))
else:
is_sd3 = False

try:
if is_diffusers_version(">=", "0.30.0"):
from diffusers import FluxPipeline

is_flux = isinstance(pipeline, FluxPipeline)
except ImportError:
flux_pipes = [FluxPipeline]

if is_diffusers_version(">=", "0.31.0"):
from diffusers import FluxImg2ImgPipeline, FluxInpaintPipeline

flux_pipes.extend([FluxPipeline, FluxImg2ImgPipeline, FluxInpaintPipeline])

if is_diffusers_version(">=", "0.32.0"):
from diffusers import FluxFillPipeline

flux_pipes.append(FluxFillPipeline)

is_flux = isinstance(pipeline, tuple(flux_pipes))
else:
is_flux = False

if not is_sd3 and not is_flux:
Expand Down
11 changes: 10 additions & 1 deletion optimum/exporters/openvino/model_configs.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,12 @@
)
from optimum.utils.normalized_config import NormalizedConfig, NormalizedTextConfig, NormalizedVisionConfig

from ...intel.utils.import_utils import _transformers_version, is_diffusers_version, is_transformers_version
from ...intel.utils.import_utils import (
_transformers_version,
is_diffusers_available,
is_diffusers_version,
is_transformers_version,
)
from .model_patcher import (
AquilaModelPatcher,
ArcticModelPatcher,
Expand Down Expand Up @@ -116,6 +121,10 @@ def init_model_configs():
"image-text-to-text"
] = TasksManager._TRANSFORMERS_TASKS_TO_MODEL_LOADERS["text-generation"]

if is_diffusers_available():
TasksManager._DIFFUSERS_TASKS_TO_MODEL_LOADERS["fill"] = "FluxFillPipeline"
TasksManager._DIFFUSERS_TASKS_TO_MODEL_MAPPINGS["fill"] = {"flux": "FluxFillPipeline"}

supported_model_types = [
"_SUPPORTED_MODEL_TYPE",
"_DIFFUSERS_SUPPORTED_MODEL_TYPE",
Expand Down
19 changes: 17 additions & 2 deletions optimum/intel/openvino/modeling_diffusion.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,11 @@
FluxImg2ImgPipeline = object
FluxInpaintPipeline = object

if is_diffusers_version(">=", "0.32.0"):
from diffusers import FluxFillPipeline
else:
FluxFillPipeline = object


DIFFUSION_MODEL_TRANSFORMER_SUBFOLDER = "transformer"
DIFFUSION_MODEL_TEXT_ENCODER_3_SUBFOLDER = "text_encoder_3"
Expand Down Expand Up @@ -1452,17 +1457,23 @@ class OVFluxPipeline(OVDiffusionPipeline, OVTextualInversionLoaderMixin, FluxPip


class OVFluxImg2ImgPipeline(OVDiffusionPipeline, OVTextualInversionLoaderMixin, FluxImg2ImgPipeline):
main_input_name = "prompt"
main_input_name = "image"
export_feature = "image-to-image"
auto_model_class = FluxImg2ImgPipeline


class OVFluxInpaintPipeline(OVDiffusionPipeline, OVTextualInversionLoaderMixin, FluxInpaintPipeline):
main_input_name = "prompt"
main_input_name = "image"
export_feature = "inpainting"
auto_model_class = FluxInpaintPipeline


class OVFluxFillPipeline(OVDiffusionPipeline, OVTextualInversionLoaderMixin, FluxFillPipeline):
main_input_name = "image"
export_feature = "inpainting"
auto_model_class = FluxFillPipeline


SUPPORTED_OV_PIPELINES = [
OVStableDiffusionPipeline,
OVStableDiffusionImg2ImgPipeline,
Expand Down Expand Up @@ -1531,6 +1542,10 @@ def _get_ov_class(pipeline_class_name: str, throw_error_if_not_exist: bool = Tru
OV_INPAINT_PIPELINES_MAPPING["flux"] = OVFluxInpaintPipeline
OV_IMAGE2IMAGE_PIPELINES_MAPPING["flux"] = OVFluxImg2ImgPipeline

if is_diffusers_version(">=", "0.32.0"):
OV_INPAINT_PIPELINES_MAPPING["flux-fill"] = OVFluxFillPipeline
SUPPORTED_OV_PIPELINES.append(OVFluxFillPipeline)

SUPPORTED_OV_PIPELINES_MAPPINGS = [
OV_TEXT2IMAGE_PIPELINES_MAPPING,
OV_IMAGE2IMAGE_PIPELINES_MAPPING,
Expand Down

0 comments on commit 58ada07

Please sign in to comment.