import torch
from diffusers import StableDiffusion3ControlNetPipeline
from diffusers.models import SD3ControlNetModel
from diffusers.utils import load_image
from diffusers.image_processor import VaeImageProcessor

device = "cuda:3"
dtype = torch.float16

# load pipeline

class SD3CannyImageProcessor(VaeImageProcessor):
    def __init__(self):
        super().__init__(do_normalize=False)
    def preprocess(self, image, **kwargs):
        image = super().preprocess(image, **kwargs)
        image = image * 255 * 0.5 + 0.5
        return image
    def postprocess(self, image, do_denormalize=True, **kwargs):
        do_denormalize = [True] * image.shape[0]
        image = super().postprocess(image, **kwargs, do_denormalize=do_denormalize)
        return image


control_image_processor = SD3CannyImageProcessor()
controlnet = SD3ControlNetModel.from_pretrained("diffusers-internal-dev/sd35-controlnet-canny-8b", torch_dtype=dtype)
pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-large",
    controlnet=controlnet,
    torch_dtype=dtype
)
pipe.to(device)
pipe.image_processor = control_image_processor

# config
control_image = load_image("https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/canny.png")
prompt =  "A Night time photo taken by Leica M11, portrait of a Japanese woman in a kimono, looking at the camera, Cherry blossoms"
max_sequence_length = 77

# to reproduce result in our example
generator = torch.Generator(device="cpu").manual_seed(23)
image = pipe(
    prompt, 
    control_image=control_image, 
    controlnet_conditioning_scale=1.0,
    guidance_scale=3.5,
    num_inference_steps=60,
    generator=generator,
    max_sequence_length=max_sequence_length,
).images[0]
image.save(f'yiyi_test_2_out_{max_sequence_length}.jpg')

import torch
from diffusers import StableDiffusion3ControlNetPipeline
from diffusers.models import SD3ControlNetModel
from diffusers.utils import load_image

device = "cuda:3"
dtype = torch.float16

# load pipeline
controlnet = SD3ControlNetModel.from_pretrained("diffusers-internal-dev/sd35-controlnet-depth-8b", torch_dtype=dtype)
pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-large",
    controlnet=controlnet,
    torch_dtype=dtype
)
pipe.to(device)

# config
control_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/marigold/marigold_einstein_lcm_depth.png")
prompt = "a photo of a man"
max_sequence_length = 77

# to reproduce result in our example
generator = torch.Generator(device="cpu").manual_seed(0)
image = pipe(
    prompt, 
    control_image=control_image, 
    controlnet_conditioning_scale=1.0,
    guidance_scale=4.5,
    num_inference_steps=40,
    generator=generator,
    max_sequence_length=max_sequence_length,
).images[0]
image.save(f'yiyi_test_1_out_{max_sequence_length}.jpg')

# test blur
import torch
from diffusers import StableDiffusion3ControlNetPipeline
from diffusers.models import SD3ControlNetModel
from diffusers.utils import load_image

device = "cuda:3"
dtype = torch.float16

# load pipeline
controlnet = SD3ControlNetModel.from_pretrained("diffusers-internal-dev/sd35-controlnet-blur-8b", torch_dtype=dtype)
pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-large",
    controlnet=controlnet,
    torch_dtype=dtype
)
pipe.to(device)

# config
control_image = load_image("https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/blur.png")
print(f" control image size: {control_image.size}")
prompt = "generated ai art, a tiny, lost rubber ducky in an action shot close-up, surfing the humongous waves, inside the tube, in the style of Kelly Slater"
max_sequence_length = 77

# to reproduce result in our example
generator = torch.Generator(device="cpu").manual_seed(0)
image = pipe(
    prompt, 
    control_image=control_image, 
    controlnet_conditioning_scale=1.0,
    guidance_scale=3.5,
    num_inference_steps=60,
    generator=generator,
    max_sequence_length=max_sequence_length,
).images[0]
image.save(f'yiyi_test_3_out_{max_sequence_length}.jpg')