import torch
from controlnet_aux import CannyDetector
from diffusers import FluxControlPipeline
from diffusers.utils import load_image
from diffusers import FluxImg2ImgPipeline


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

pipe.load_lora_weights(
    "black-forest-labs/FLUX.1-Canny-dev-lora"
)

prompt = "A robot made of exotic candies and chocolates of different kinds. The background is filled with confetti and celebratory gifts."
control_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/robot.png")

processor = CannyDetector()
control_image = processor(control_image, low_threshold=50, high_threshold=200, detect_resolution=1024, image_resolution=1024)

image = pipe(
    prompt=prompt,
    control_image=control_image,
    num_inference_steps=50,
    guidance_scale=30.0,
).images[0]

# pipe = FluxImg2ImgPipeline.from_pretrained(
#     "black-forest-labs/FLUX.1-dev", 
#     torch_dtype=torch.bfloat16
# )

pipe = FluxImg2ImgPipeline.from_pipe(
    pipe,
    torch_dtype=torch.bfloat16
)

pipe = pipe.to("cuda")

init_image = load_image("https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg").resize((1024, 1024))
prompt = "cat wizard, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney, 8k"
image = pipe(
    prompt=prompt, 
    image=init_image, 
    num_inference_steps=28, 
    strength=0.5, 
    guidance_scale=2.5
).images[0]

pipe = FluxImg2ImgPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev", 
     torch_dtype=torch.bfloat16
)