Flux with Remote Encode

[hlky]

What does this PR do?

#11069

import torch
from diffusers import FluxImg2ImgPipeline
from diffusers.utils import load_image
from diffusers.utils.constants import ENCODE_ENDPOINT_FLUX, DECODE_ENDPOINT_FLUX
from diffusers.utils.remote_utils import remote_decode, remote_encode
device = "cuda"
pipe = FluxImg2ImgPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", vae=None, torch_dtype=torch.bfloat16)
pipe = pipe.to(device)
url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
init_image = load_image(url).resize((1024, 1024))

init_image = remote_encode(
    endpoint=ENCODE_ENDPOINT_FLUX,
    image=init_image,
    scaling_factor=0.3611,
    shift_factor=0.1159,
)

prompt = "cat wizard, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney, 8k"
latents = pipe(
    prompt=prompt, image=init_image, strength=0.95, output_type="latent"
).images

image = remote_decode(
    endpoint=DECODE_ENDPOINT_FLUX,
    tensor=latents,
    scaling_factor=0.3611,
    shift_factor=0.1159,
    partial_postprocess=True,
    output_type="pt",
    return_type="pil",
    height=1024,
    width=1024,
)
image.save("wizard.png")

import torch
from diffusers import FluxInpaintPipeline
from diffusers.utils import load_image
from diffusers.utils.constants import ENCODE_ENDPOINT_FLUX, DECODE_ENDPOINT_FLUX
from diffusers.utils.remote_utils import remote_decode, remote_encode

pipe = FluxInpaintPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", vae=None, torch_dtype=torch.bfloat16
).to("cuda")

prompt = "Face of a yellow cat, high resolution, sitting on a park bench"

img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"

init_image = load_image(img_url)
mask = load_image(mask_url)

mask_condition = pipe.mask_processor.preprocess(mask, height=1024, width=1024)
init_image = pipe.image_processor.preprocess(init_image, height=1024, width=1024)
masked_image = init_image * (mask_condition < 0.5)

init_image = remote_encode(
    endpoint=ENCODE_ENDPOINT_FLUX,
    image=init_image,
    scaling_factor=0.3611,
    shift_factor=0.1159,
)

masked_image = remote_encode(
    endpoint=ENCODE_ENDPOINT_FLUX,
    image=masked_image,
    scaling_factor=0.3611,
    shift_factor=0.1159,
)

latents = pipe(
    prompt=prompt,
    image=init_image,
    mask_image=mask,
    masked_image_latents=masked_image,
    output_type="latent",
).images

image = remote_decode(
    endpoint=DECODE_ENDPOINT_FLUX,
    tensor=latents,
    scaling_factor=0.3611,
    shift_factor=0.1159,
    partial_postprocess=True,
    output_type="pt",
    return_type="pil",
    height=1024,
    width=1024,
)
image.save("cat.png")

• Img2Img
• Inpaint

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[hlky]

cc @yiyixuxu Should we make output_type="latent" unpack? The packed latents cannot be passed to latents as input and we need to pass height and width to remote_encode for packed latents.

diffusers/src/diffusers/pipelines/flux/pipeline_flux_img2img.py

Lines 1013 to 1020 in 33d10af

	if output_type == "latent":
	image = latents
	else:
	latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
	latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
	image = self.vae.decode(latents, return_dict=False)[0]
	image = self.image_processor.postprocess(image, output_type=output_type)

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[hlky]

cc @vladmandic Here is Flux img2img

[yiyixuxu]

awesome!

[hlky]

cc @vladmandic Inpaint done

[hlky]

Merging to not unblock downstream.

Gentle ping @yiyixuxu on #11091 (comment) for a follow up PR.