inference_flux2.py

from pipelines.pipeline_flux2 import Flux2MultiDiffusionPipeline
import torch
from diffusers.utils import load_image
from PIL import Image
Image.MAX_IMAGE_PIXELS = None
from pathlib import Path

# Set the base model
base_model = "black-forest-labs/FLUX.2-dev"

inference_dtype = torch.bfloat16

base = Flux2MultiDiffusionPipeline.from_pretrained(
    base_model,
    torch_dtype=inference_dtype,
)

# Enable memory optimizations. You can choose between CPU offloading or sequential CPU offloading.
#base.enable_model_cpu_offload(0)
base.enable_sequential_cpu_offload(gpu_id=0)

# Instance prompt for testing. You can change this to test different prompts.
instanceprompt = ("A breathtaking, ultra-wide panoramic landscape of jagged obsidian mountain peaks at twilight. The valleys are filled with flowing rivers of glowing violet bioluminescent mist. Above, a sprawling cosmic sky features a massive, detailed nebula in shades of deep teal and gold, with scattered crystalline stars. Photorealistic, cinematic lighting, 8k resolution, highly detailed rock textures, atmospheric depth, ethereal glow, wide-angle lens.")

# Directory path for output images. The directory will be created if it doesn't exist. 
output_path = Path("./outputs")
output_path.mkdir(parents=True, exist_ok=True)
img_name = "res_flux2_MultiDiffusion"

# Multidiffusion image size
TARGET_HEIGHT = 3072 # 3072
TARGET_WIDTH = 3072 # 3072

# Generate the image using the FLUX.2-dev pipeline with the specified parameters.
image = base(
    prompt = instanceprompt,
    height = 1024, # inference height
    width = 1024, # inference width
    total_height = TARGET_HEIGHT, # Multidiffusion target height
    total_width = TARGET_WIDTH, # Multidiffusion target width
    num_inference_steps = 20, # reduce for faster testing, increase for better quality
    guidance_scale = 4.0, # helps FLUX.2-dev follow the prompt
    view_batch_size = 1, 
    stride_start= 512, # Multidiffusion stride start (minimum stride for the first denoising step)
    stride_end= 512 + 256, # Multidiffusion stride end (maximum stride for the last denoising step). The stride will change linearly from stride_start to stride_end across the denoising steps.
    return_dict = False
)[0][0]

# Save the generated image to the specified output path. The image will be saved in PNG format with the name defined in img_name.
image.save(output_path / f"{img_name}.png")