from pred_color import *
import gradio as gr

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDPMScheduler,
    StableDiffusionControlNetPipeline,
    UNet2DConditionModel,
    UniPCMultistepScheduler,
)
import torch
from diffusers.utils import load_image

controlnet_model_name_or_path = "svjack/ControlNet-Face-Zh"
controlnet = ControlNetModel.from_pretrained(controlnet_model_name_or_path)
#controlnet = controlnet.to("cuda")

base_model_path = "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1"
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    base_model_path, controlnet=controlnet,
    #torch_dtype=torch.float16
)

# speed up diffusion process with faster scheduler and memory optimization
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.safety_checker = None
#pipe.enable_model_cpu_offload()
#pipe = pipe.to("cuda")

if torch.cuda.is_available():
    pipe = pipe.to("cuda")
else:
    #pipe.enable_model_cpu_offload()
    pass

example_sample = [
    ["Protector_Cromwell_style.png", "戴帽子穿灰色衣服的男子"]
]

from PIL import Image
def pred_func(image, prompt):
    #out = single_pred_features(image)
    features ,face_features = single_pred_features(image)
    req_img = produce_center_crop_image(features ,face_features)
    out = {}
    out["spiga_seg"] = req_img
    if type(out) == type({}):
        #return out["spiga_seg"]
        control_image = out["spiga_seg"]
        if type(image) == type("") and os.path.exists(image):
            image = Image.open(image).convert("RGB")
        elif hasattr(image, "shape"):
            image = Image.fromarray(image).convert("RGB")
        else:
            image = image.convert("RGB")
        image = image.resize((512, 512))

        generator = torch.manual_seed(0)
        image = pipe(
             prompt, num_inference_steps=50,
             generator=generator, image=control_image
        ).images[0]
        return control_image ,image


gr=gr.Interface(fn=pred_func, inputs=['image','text'],
outputs=[gr.Image(label='output').style(height=512),
gr.Image(label='output').style(height=512)],
examples=example_sample if example_sample else None,
cache_examples = False
)
gr.launch(share=False)