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


def run_pipeline(test="offload"):
    print(" ")
    print(f"testing pipeline with {test}")

    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()

    print("Loading pipeline")
    controlnet = SD3ControlNetModel.from_pretrained("stabilityai/stable-diffusion-3.5-large-controlnet-depth", torch_dtype=torch.float16)
    pipe = StableDiffusion3ControlNetPipeline.from_pretrained(
        "stabilityai/stable-diffusion-3.5-large",
        controlnet=controlnet,
        torch_dtype=torch.float16,
        device_map="balanced" if test=="device_map" else None,
    )

    if test=="offload":
        pipe.enable_model_cpu_offload()
    elif test == "cuda":
        pipe.to("cuda")

    peak_mem = torch.cuda.max_memory_allocated() / 1024**3
    print(f"Peak memory after loading: {peak_mem:.2f} GB")
    

    control_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/marigold/marigold_einstein_lcm_depth.png")
    generator = torch.Generator(device="cpu").manual_seed(0)
 
    print("Running pipeline")
    try:
        image = pipe(
            prompt = "a photo of a man", 
            control_image=control_image, 
            guidance_scale=4.5,
            num_inference_steps=40,
            generator=generator,
            max_sequence_length=77,
        ).images[0]
        final_peak_mem = torch.cuda.max_memory_allocated() / 1024**3
        print(f"Peak memory after inference: {final_peak_mem:.2f} GB")
        image.save(f'yiyi_test_1_{test}_out.png')
    except Exception as e:
        print(e)
        for n,m in pipe.components.items():
            if isinstance(m, torch.nn.Module):
                print(f" -{n}: {m.device}")
        print(f" ")
    


run_pipeline(test="offload")
run_pipeline(test="cuda")
# run_pipeline(test="device_map")