import torch
from diffusers import FluxControlInpaintPipeline
from diffusers.models.transformers import FluxTransformer2DModel
from transformers import T5EncoderModel
from diffusers.utils import load_image, make_image_grid
from image_gen_aux import DepthPreprocessor
from PIL import Image
import numpy as np

pipe = FluxControlInpaintPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Depth-dev",
    torch_dtype=torch.bfloat16,
)
# GPU optimization code...
pipe.to("cuda")

prompt = "a blue robot singing opera with human-like expressions"
image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/robot.png")

# Create mask for the robot's head
head_mask = np.zeros_like(image)
head_mask[65:580,300:642] = 255
mask_image = Image.fromarray(head_mask)

# Process depth map
processor = DepthPreprocessor.from_pretrained("LiheYoung/depth-anything-large-hf")
control_image = processor(image)[0].convert("RGB")

output = pipe(
    prompt=prompt,
    image=image,
    control_image=control_image,
    mask_image=mask_image,
    num_inference_steps=30,
    strength=0.9,
    guidance_scale=10.0,
    generator=torch.Generator().manual_seed(42),
).images[0]

import torch
import numpy as np
from PIL import Image
from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel
from diffusers.utils import make_image_grid
import controlnet_hinter
import cv2

# Setup model
controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-depth", torch_dtype=torch.float16)
pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
    "...inpaint_model...",
    controlnet=controlnet,
    torch_dtype=torch.float16
)
pipe.to("cuda")

# Load images
init_image = Image.open('image.jpg')
mask_img = Image.open('mask.png')

# Prepare mask
mask_array = np.array(mask_img) > 0
mask_array = cv2.resize(mask_array.astype(np.uint8), 
                     (init_image.size[0], init_image.size[1]), 
                     interpolation=cv2.INTER_NEAREST).astype(bool)
mask = Image.fromarray(mask_array.astype(np.uint8) * 255)

# Generate depth map
control_image = controlnet_hinter.hint_depth(init_image)
control_image = control_image.resize(init_image.size)

# Generation configuration
generator = torch.Generator("cuda").manual_seed(42)
config = {
    "negative_prompt": "bad quality, worst quality",
    "num_inference_steps": 30,
    "guidance_scale": 7.5,
    "strength": 0.7,
    "controlnet_conditioning_scale": 0.6,
    "control_guidance_start": 0.6,
    "control_guidance_end": 0.8
}

# Generate image
output = pipe(
    prompt="Elegant blonde woman displaying refined style...",
    image=init_image,
    mask_image=mask,
    control_image=control_image,
    # other parameters...
).images[0]