# app.py
import torch
import gradio as gr
from diffusers import DiffusionPipeline

print("Loading pipeline...")

device = "cuda" if torch.cuda.is_available() else "cpu"

dtype = torch.float16 if device == "cuda" else torch.float32

pipe = DiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5",
    torch_dtype=torch.float16 if device=="cuda" else torch.float32,
    cache_dir="/tmp/huggingface",
    use_safetensors=True,
    safety_checker=None
)

# pipe.safety_checker = None

# device = "cuda" if torch.cuda.is_available() else "cpu"

pipe.to(device)

pipe.enable_attention_slicing()
# pipe.enable_model_cpu_offload()


# pipe.unet.load_attn_procs(
# "./pytorch_custom_diffusion_weights.bin"
# )

# attn_path = "./pytorch_custom_diffusion_weights.bin"
# state_dict = torch.load(attn_path, map_location="cpu")
# pipe.unet.load_attn_procs(
# state_dict
# ) 

import os
print(os.path.getsize("pytorch_custom_diffusion_weights.bin"))

pipe.unet.load_attn_procs(
    "./pytorch_custom_diffusion_weights.bin",
    weight_name="pytorch_custom_diffusion_weights.bin"
)

print("Pipeline loaded")

# def generate(prompt, steps, guidance):

#     print("Generating...")
    
#     image = pipe(
#         prompt,
#         num_inference_steps=steps,
#         guidance_scale=guidance,
#         eta=1
#     ).images[0]

#     print("Done")
    
#     return image

def generate(prompt, steps, guidance):

    print("Generating...")

    result = pipe(
        prompt,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance)
    )

    print("RESULT:", result)

    image = result.images[0]

    return image


demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(
            label="Prompt",
            value="A <new1> reference. New Year image with a rabbit as the main element"
        ),
        gr.Slider(10, 320, value=100, label="Steps"),
        gr.Slider(1, 18, value=6, label="Guidance"),
    ],
    outputs=gr.Image(),
    title="Fine-tuning style diffusion Demo"
)

demo.launch(
    server_name="0.0.0.0",
    server_port=7860
)