Spaces:
Runtime error
Runtime error
File size: 5,047 Bytes
86a2bd4 bb21a7c 94b2c21 bb21a7c 4564155 bb21a7c e5d785e bb21a7c 94b2c21 bb21a7c 94b2c21 f856c81 94b2c21 bb21a7c 94b2c21 bb21a7c 94b2c21 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
from controlnet_aux import OpenposeDetector
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers import UniPCMultistepScheduler
import gradio as gr
import torch
import base64
from io import BytesIO
from PIL import Image
# live conditioning
canvas_html = "<pose-canvas id='canvas-root' style='display:flex;max-width: 500px;margin: 0 auto;'></pose-canvas>"
load_js = """
async () => {
const url = "https://huggingface.co/datasets/radames/gradio-components/raw/main/pose-gradio.js"
fetch(url)
.then(res => res.text())
.then(text => {
const script = document.createElement('script');
script.type = "module"
script.src = URL.createObjectURL(new Blob([text], { type: 'application/javascript' }));
document.head.appendChild(script);
});
}
"""
get_js_image = """
async (image_in_img, prompt, image_file_live_opt, live_conditioning) => {
const canvasEl = document.getElementById("canvas-root");
const data = canvasEl? canvasEl._data : null;
return [image_in_img, prompt, image_file_live_opt, data]
}
"""
# Constants
low_threshold = 100
high_threshold = 200
# Models
pose_model = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
controlnet = ControlNetModel.from_pretrained(
"lllyasviel/sd-controlnet-openpose", torch_dtype=torch.float16
)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
# This command loads the individual model components on GPU on-demand. So, we don't
# need to explicitly call pipe.to("cuda").
pipe.enable_model_cpu_offload()
# xformers
pipe.enable_xformers_memory_efficient_attention()
# Generator seed,
generator = torch.manual_seed(0)
def get_pose(image):
return pose_model(image)
def generate_images(image, prompt, image_file_live_opt='file', live_conditioning=None):
if image is None and 'image' not in live_conditioning:
raise gr.Error("Please provide an image")
try:
if image_file_live_opt == 'file':
pose = get_pose(image)
elif image_file_live_opt == 'webcam':
base64_img = live_conditioning['image']
image_data = base64.b64decode(base64_img.split(',')[1])
pose = Image.open(BytesIO(image_data)).convert(
'RGB').resize((512, 512))
output = pipe(
prompt,
pose,
generator=generator,
num_images_per_prompt=3,
num_inference_steps=20,
)
all_outputs = []
all_outputs.append(pose)
for image in output.images:
all_outputs.append(image)
return all_outputs
except Exception as e:
raise gr.Error(str(e))
def toggle(choice):
if choice == "file":
return gr.update(visible=True, value=None), gr.update(visible=False, value=None)
elif choice == "webcam":
return gr.update(visible=False, value=None), gr.update(visible=True, value=canvas_html)
with gr.Blocks() as blocks:
gr.Markdown("""
## Generate controlled outputs with ControlNet and Stable Diffusion
This Space uses pose estimated lines as the additional conditioning
[Check out our blog to see how this was done (and train your own controlnet)](https://huggingface.co/blog/train-your-controlnet)
""")
with gr.Row():
live_conditioning = gr.JSON(value={}, visible=False)
with gr.Column():
image_file_live_opt = gr.Radio(["file", "webcam"], value="file",
label="How would you like to upload your image?")
image_in_img = gr.Image(source="upload", visible=True, type="pil")
canvas = gr.HTML(None, elem_id="canvas_html", visible=False)
image_file_live_opt.change(fn=toggle,
inputs=[image_file_live_opt],
outputs=[image_in_img, canvas],
queue=False)
prompt = gr.Textbox(
label="Enter your prompt",
max_lines=1,
placeholder="best quality, extremely detailed",
)
run_button = gr.Button("Generate")
with gr.Column():
gallery = gr.Gallery().style(grid=[2], height="auto")
run_button.click(fn=generate_images,
inputs=[image_in_img, prompt,
image_file_live_opt, live_conditioning],
outputs=[gallery],
_js=get_js_image)
blocks.load(None, None, None, _js=load_js)
gr.Examples(fn=generate_images,
examples=[
["./yoga1.jpeg",
"best quality, extremely detailed"]
],
inputs=[image_in_img, prompt],
outputs=[gallery],
cache_examples=True)
blocks.launch(debug=True)
|