import gradio as gr
from controlnet_aux import OpenposeDetector
from PIL import Image
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
import torch
from controlnet_aux import OpenposeDetector
from diffusers.utils import load_image

#Models
openpose = OpenposeDetector.from_pretrained('lllyasviel/ControlNet')
controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-openpose", torch_dtype=torch.float16)
pipe = StableDiffusionControlNetPipeline.from_pretrained("/content/drive/MyDrive/hack/800", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16)

#optimizations
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_xformers_memory_efficient_attention()
pipe.enable_model_cpu_offload()

def generate(image,prompt):
    image = openpose(image)
    print(image)
    image = pipe(prompt, image, num_inference_steps=20).images[0]
    return image

gr.Interface(fn=generate, inputs=["image","text"], outputs="image").launch(share=True)