import torch
import os
from huggingface_hub import HfApi
from pathlib import Path
from diffusers.utils import load_image
from PIL import Image
import numpy as np
from controlnet_aux import OpenposeDetector

from diffusers import (
    ControlNetModel,
    StableDiffusionControlNetPipeline,
    UniPCMultistepScheduler,
)

checkpoint = "lllyasviel/control_v11p_sd15_openpose"


processor = OpenposeDetector.from_pretrained('lllyasviel/ControlNet')

controlnet = ControlNetModel.from_pretrained(checkpoint, torch_dtype=torch.float16)

pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
)

pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()

def openposeImage(image, prompt):
  prompt += "oil painting, yoji shinakawa, studio gainax, y2k design, dramatic lighting"
  control_image = processor(image, hand_and_face=True)
  generator = torch.manual_seed(0)
  image = pipe(prompt, num_inference_steps=30, generator=generator, image=control_image).images[0]

  return image

import gradio as gr
gr.Interface(fn=openposeImage, 
             inputs=[gr.Image(shape=(512, 512)), gr.inputs.Textbox(label="Text Prompt")],
             outputs=gr.Image(shape=(512, 512))).launch(inline = False)