save issue

it works if you trust remote source and ummm...load the 'custom pipeline' accurately:

import torch
from diffusers import AutoPipelineForText2Image
from diffusers.utils import load_image
from transformers import T5EncoderModel, TorchAoConfig
from diffusers import FluxTransformer2DModel

Model and image paths

name_or_path = "ostris/Flex.2-preview"
inpaint_image = load_image("https://ostris.com/wp-content/uploads/2025/04/dog.jpg")
inpaint_mask = load_image("https://ostris.com/wp-content/uploads/2025/04/dog_mask.jpg")
control_image = load_image("https://ostris.com/wp-content/uploads/2025/04/dog_depth.jpg")

Quantization and dtype

dtype = torch.bfloat16
quant_config = TorchAoConfig("int8_weight_only")

Load text encoder and transformer

text_encoder_2 = T5EncoderModel.from_pretrained(
name_or_path, subfolder="text_encoder_2", torch_dtype=dtype, quantization_config=quant_config
).to("cuda")
transformer = FluxTransformer2DModel.from_pretrained(
name_or_path, subfolder="transformer", torch_dtype=dtype, quantization_config=quant_config
).to("cuda")

Load pipeline with explicit custom_pipeline path

custom_pipeline_path = "C:\Users\UPDATE_THIS_PATH\.cache\huggingface\hub\models--ostris--Flex.2-preview\snapshots\04512ad16ef6072afba83ec30fa2f1ea95cd5636\pipeline.py"
pipe = AutoPipelineForText2Image.from_pretrained(
name_or_path,
transformer=transformer,
text_encoder_2=text_encoder_2,
custom_pipeline=custom_pipeline_path,
torch_dtype=dtype,
trust_remote_code=True
).to("cuda")

Generate image

image = pipe(
prompt="A white friendly robotic dog sitting on a bench",
inpaint_image=inpaint_image,
inpaint_mask=inpaint_mask,
control_image=control_image,
control_strength=0.5,
control_stop=0.33,
height=1024,
width=1024,
guidance_scale=3.5,
num_inference_steps=50,
generator=torch.Generator("cpu").manual_seed(42)
).images[0]

Save image

image.save("robot_dog.png")