import torch
from controlnet_aux import LineartDetector
from diffusers import ControlNetModel,UniPCMultistepScheduler,FluxPipeline

device= "cuda" if torch.cuda.is_available() else "cpu"
print("Using device for I2I_2:", device)
processor = LineartDetector.from_pretrained("lllyasviel/Annotators")

checkpoint = "ControlNet-1-1-preview/control_v11p_sd15_lineart"
controlnet = ControlNetModel.from_pretrained(checkpoint, torch_dtype=torch.float16).to(device)
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", controlnet=controlnet, torch_dtype=torch.float16
).to(device)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
# pipe.enable_model_cpu_offload()

def I2I_2(image, prompt,size,num_inference_steps):
    image.resize((size))
    image=processor(image)
    generator = torch.Generator(device=device).manual_seed(0)
    image = pipe(prompt, num_inference_steps=num_inference_steps, generator=generator, image=image).images[0]
    return image