import gradio as gr
import torch
from annotator.util import resize_image, HWC3
from cldm.model import create_model, load_state_dict
from cldm.ddim_hacked import DDIMSampler
from huggingface_hub import hf_hub_download

# Initialize the model and other components
# config = "./models/cldm_v21_512_latctrl_coltrans.yaml'"
model = create_model('./models/cldm_v21_512_latctrl_coltrans.yaml').cpu()
ckpt = hf_hub_download(repo_id="xywwww/scene_diffusion", filename="checkpoints/epoch=25-step=112553.ckpt")
print(ckpt)
model.load_state_dict(load_state_dict(ckpt), strict=False)
# model = load_model_checkpoint(model, ckpt)
ddim_sampler = DDIMSampler(model)

def process(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode, strength, scale, seed, eta, low_threshold, high_threshold):
    with torch.no_grad():
        img = resize_image(HWC3(input_image), image_resolution)
        H, W, C = img.shape
        # detected_map = apply_canny(img, low_threshold, high_threshold)
        # detected_map = HWC3(detected_map)
        # Add the rest of the processing logic here

def create_demo(process):
    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                input_image = gr.Image()
                prompt = gr.Textbox(label="Prompt", submit_btn=True)
                a_prompt = gr.Textbox(label="Additional Prompt")
                n_prompt = gr.Textbox(label="Negative Prompt")
                with gr.Accordion("Advanced options", open=False):
                    num_samples = gr.Slider(label="Number of images", minimum=1, maximum=10, value=1, step=1)
                    image_resolution = gr.Slider(label="Image resolution", minimum=256, maximum=1024, value=512, step=256)
                    ddim_steps = gr.Slider(label="DDIM Steps", minimum=1, maximum=100, value=50, step=1)
                    guess_mode = gr.Checkbox(label="Guess Mode")
                    strength = gr.Slider(label="Strength", minimum=0.0, maximum=1.0, value=0.5, step=0.1)
                    scale = gr.Slider(label="Scale", minimum=0.1, maximum=30.0, value=10.0, step=0.1)
                    seed = gr.Slider(label="Seed", minimum=0, maximum=10000, value=42, step=1)
                    eta = gr.Slider(label="ETA", minimum=0.0, maximum=1.0, value=0.0, step=0.1)
                    low_threshold = gr.Slider(label="Canny Low Threshold", minimum=1, maximum=255, value=100, step=1)
                    high_threshold = gr.Slider(label="Canny High Threshold", minimum=1, maximum=255, value=200, step=1)
                submit = gr.Button("Generate")
            with gr.Column():
                output_image = gr.Image()
        submit.click(fn=process, inputs=[input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode, strength, scale, seed, eta, low_threshold, high_threshold], outputs=output_image)
    return demo

demo = create_demo(process)
demo.launch()