app.py

import torch
from diffusers.utils import load_image
from diffusers import FluxControlNetPipeline, FluxControlNetModel
import os

# Clear unnecessary memory
torch.cuda.empty_cache()

# Set the environment variable to handle memory fragmentation
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

base_model = 'black-forest-labs/FLUX.1-dev'
controlnet_model = 'InstantX/FLUX.1-dev-Controlnet-Union'

# Use mixed precision with float16
controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16)
pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.float16)

# Offload some parts to CPU
pipe.enable_model_cpu_offload()

pipe.to("cuda")

control_image_canny = load_image("https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union-alpha/resolve/main/images/canny.jpg")
controlnet_conditioning_scale = 0.5
control_mode = 0

# Reduce the image resolution if needed to fit into memory
width, height = control_image_canny.size
width = width // 2
height = height // 2

prompt = 'A bohemian-style female travel blogger with sun-kissed skin and messy beach waves.'

image = pipe(
    prompt, 
    control_image=control_image_canny,
    control_mode=control_mode,
    width=width,
    height=height,
    controlnet_conditioning_scale=controlnet_conditioning_scale,
    num_inference_steps=24, 
    guidance_scale=3.5,
).images[0]
image.save("image.jpg")

# Empty cache after the operation to free up memory
torch.cuda.empty_cache()