# import gradio as gr
# import torch
# from PIL import Image
# from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
# import logging

# # 配置日志记录
# logging.basicConfig(level=logging.DEBUG)

# # 加载模型和ControlNet
# model_id = "CompVis/stable-diffusion-v1-4"
# controlnet_id = "lllyasviel/sd-controlnet-canny"

# pipe = StableDiffusionControlNetPipeline.from_pretrained(model_id, controlnet=ControlNetModel.from_pretrained(controlnet_id))
# pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")

# def is_valid_pil_image(image):
#     return isinstance(image, Image.Image)

# def generate_image(control_image, additional_prompt, reference_image1, reference_image2):
#     try:
#         logging.debug("Starting image generation process...")
        
#         # 打印传入的图像对象类型
#         logging.debug(f"Control Image type: {type(control_image)}")
#         logging.debug(f"Reference Image 1 type: {type(reference_image1)}")
#         logging.debug(f"Reference Image 2 type: {type(reference_image2)}")
        
#         # 检查输入图像是否为有效的 PIL 图像
#         if not all(is_valid_pil_image(img) for img in [control_image, reference_image1, reference_image2]):
#             error_message = "Error: One or more input images are not valid PIL images."
#             logging.error(error_message)
#             return error_message

#         # 第一步：使用 canny 模型生成初始图像
#         initial_image = pipe(prompt=additional_prompt, control_image=control_image).images[0]
#         if not is_valid_pil_image(initial_image):
#             error_message = "Error in Step 1: Generated image is not a valid PIL image."
#             logging.error(error_message)
#             return error_message

#         # 第二步：使用第一张参考图像生成图像
#         step2_image = pipe(prompt=additional_prompt, control_image=reference_image1).images[0]
#         if step2_image is None:
#             error_message = "Error in Step 2: Generated image is None (NoneType)."
#             logging.error(error_message)
#             return error_message
#         if not is_valid_pil_image(step2_image):
#             error_message = "Error in Step 2: Generated image is not a valid PIL image."
#             logging.error(error_message)
#             return error_message

#         # 第三步：使用第二张参考图像生成图像
#         final_image = pipe(prompt=additional_prompt, control_image=reference_image2).images[0]
#         if final_image is None:
#             error_message = "Error in Step 3: Generated image is None (NoneType)."
#             logging.error(error_message)
#             return error_message
#         if not is_valid_pil_image(final_image):
#             error_message = "Error in Step 3: Generated image is not a valid PIL image."
#             logging.error(error_message)
#             return error_message

#         logging.debug("Image generation process completed successfully.")
#         return final_image
#     except Exception as e:
#         error_message = f"Error: {str(e)}"
#         logging.error(error_message)
#         return error_message

# # 定义Gradio接口
# interface = gr.Interface(
#     fn=generate_image,
#     inputs=[
#         gr.components.Image(type="pil", label="Control Image"),
#         gr.components.Textbox(label="Additional Prompt"),
#         gr.components.Image(type="pil", label="Reference Image 1"),
#         gr.components.Image(type="pil", label="Reference Image 2")
#     ],
#     outputs=gr.components.Image(type="pil", label="Generated Image"),
#     title="Stable Diffusion with Multi-Step ControlNet",
#     description="Generate images using a multi-step process with Stable Diffusion and ControlNet."
# )

# if __name__ == "__main__":
#     interface.launch()

# import gradio as gr
# import torch
# from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
# from diffusers.utils import load_image
# import cv2
# import numpy as np
# import logging
# from PIL import Image

# # 配置日志记录
# logging.basicConfig(level=logging.DEBUG)

# # 加载模型
# # model_id = "runwayml/stable-diffusion-v1-5"
# model_id = "digiplay/Cetus-Mix-Codaedition_diffusers"
# controlnet_id = "lllyasviel/sd-controlnet-canny"
# device = "cpu"

# # 加载 ControlNet 模型
# controlnet = ControlNetModel.from_pretrained(controlnet_id, torch_dtype=torch.float32)

# # 加载 Stable Diffusion 管道
# pipe = StableDiffusionControlNetPipeline.from_pretrained(
#     model_id,
#     controlnet=controlnet,
#     use_safetensors=True,
#     torch_dtype=torch.float32
# )
# pipe = pipe.to(device)

# def canny_edge_detection(image):
#     image = np.array(image)
#     edges = cv2.Canny(image, 100, 200)
#     return Image.fromarray(edges)

# def generate_image(prompt, input_image):
#     try:
#         logging.debug("Starting image generation process...")

#         # 进行 Canny 边缘检测
#         canny_image = canny_edge_detection(input_image)

#         # 使用 ControlNet 和 Stable Diffusion 模型生成图像
#         with torch.no_grad():
#             generated_images = pipe(
#                 prompt=prompt,
#                 image=canny_image,
#                 height=500,
#                 width=750,
#                 num_inference_steps=20
#             ).images

#         if not generated_images or not isinstance(generated_images[0], Image.Image):
#             error_message = "Error: Generated image is not valid."
#             logging.error(error_message)
#             return error_message

#         final_image = generated_images[0]
#         logging.debug("Image generation process completed successfully.")
#         return final_image
#     except Exception as e:
#         error_message = f"Error: {str(e)}"
#         logging.error(error_message)
#         return error_message

# # 定义 Gradio 接口
# interface = gr.Interface(
#     fn=generate_image,
#     inputs=[
#         gr.components.Textbox(label="Prompt"),
#         gr.components.Image(type="pil", label="Input Image")
#     ],
#     outputs=gr.components.Image(type="pil", label="Generated Image"),
#     title="Stable Diffusion with ControlNet Canny",
#     description="Generate images using a prompt and an input image with Stable Diffusion and ControlNet Canny on CPU."
# )

# if __name__ == "__main__":
#     interface.launch()

import gradio as gr
import torch
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
import cv2
import numpy as np
import logging
from PIL import Image

# 配置日志记录
logging.basicConfig(level=logging.DEBUG)

# 设置设备
device = "cpu"
torch_dtype = torch.float32

# 加载模型
model_id = "digiplay/Cetus-Mix-Codaedition_diffusers"
canny_controlnet_id = "lllyasviel/sd-controlnet-canny"
reference_controlnet_id = "lllyasviel/control_v11f1e_sd15_tile"

# 加载 ControlNet 模型
canny_controlnet = ControlNetModel.from_pretrained(canny_controlnet_id, torch_dtype=torch_dtype)
reference_controlnet = ControlNetModel.from_pretrained(reference_controlnet_id, torch_dtype=torch_dtype)

# 加载 Stable Diffusion 管道
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    model_id,
    controlnet=[canny_controlnet, reference_controlnet],
    use_safetensors=True,
    torch_dtype=torch_dtype
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to(device)

def canny_edge_detection(image):
    image = np.array(image)
    edges = cv2.Canny(image, 100, 200)
    return Image.fromarray(edges)

def generate_image(prompt, input_image):
    try:
        logging.debug("Starting image generation process...")

        # 进行 Canny 边缘检测
        canny_image = canny_edge_detection(input_image)

        # 使用 MultiControlNet 和 Stable Diffusion 模型生成图像
        with torch.no_grad():
            generated_images = pipe(
                prompt,
                image=[canny_image, input_image],  # Canny 图像和原始输入图像
                controlnet_conditioning_scale=[1.0, 0.75],  # Canny 和 reference 的权重
                guess_mode=False,  # 对 reference 使用 False
                num_inference_steps=20,
                height=500,
                width=750,
            ).images

        if not generated_images or not isinstance(generated_images[0], Image.Image):
            error_message = "Error: Generated image is not valid."
            logging.error(error_message)
            return error_message

        final_image = generated_images[0]
        logging.debug("Image generation process completed successfully.")
        return final_image
    except Exception as e:
        error_message = f"Error: {str(e)}"
        logging.error(error_message)
        return error_message

# 定义 Gradio 接口
interface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.components.Textbox(label="Prompt"),
        gr.components.Image(type="pil", label="Input Image")
    ],
    outputs=gr.components.Image(type="pil", label="Generated Image"),
    title="Stable Diffusion with MultiControlNet (Canny and Reference)",
    description="Generate images using a prompt and an input image with Stable Diffusion and MultiControlNet (Canny and Reference) on CPU."
)

if __name__ == "__main__":
    interface.launch()