|
import gradio as gr |
|
import jax |
|
import jax.numpy as jnp |
|
import numpy as np |
|
from flax.jax_utils import replicate |
|
from flax.training.common_utils import shard |
|
from PIL import Image |
|
from diffusers import FlaxStableDiffusionControlNetPipeline, FlaxControlNetModel |
|
import cv2 |
|
|
|
with open("test.html") as f: |
|
lines = f.readlines() |
|
|
|
def create_key(seed=0): |
|
return jax.random.PRNGKey(seed) |
|
|
|
|
|
|
|
|
|
|
|
def wandb_report(url): |
|
iframe = f'<iframe src ={url} style="border:none;height:1024px;width:100%"/frame>' |
|
return gr.HTML(iframe) |
|
|
|
report_url = 'https://wandb.ai/john-fozard/dog-cat-pose/runs/kmwcvae5' |
|
control_img = 'myimage.jpg' |
|
|
|
controlnet, controlnet_params = FlaxControlNetModel.from_pretrained( |
|
"JFoz/dog-cat-pose", dtype=jnp.bfloat16 |
|
) |
|
pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained( |
|
"runwayml/stable-diffusion-v1-5", controlnet=controlnet, revision="flax", dtype=jnp.bfloat16 |
|
) |
|
|
|
def infer(prompts, negative_prompts, image): |
|
|
|
params["controlnet"] = controlnet_params |
|
|
|
num_samples = 1 |
|
rng = create_key(0) |
|
rng = jax.random.split(rng, jax.device_count()) |
|
image = Image.fromarray(image) |
|
|
|
prompt_ids = pipe.prepare_text_inputs([prompts] * num_samples) |
|
negative_prompt_ids = pipe.prepare_text_inputs([negative_prompts] * num_samples) |
|
processed_image = pipe.prepare_image_inputs([image] * num_samples) |
|
|
|
p_params = replicate(params) |
|
prompt_ids = shard(prompt_ids) |
|
negative_prompt_ids = shard(negative_prompt_ids) |
|
processed_image = shard(processed_image) |
|
|
|
output = pipe( |
|
prompt_ids=prompt_ids, |
|
image=processed_image, |
|
params=p_params, |
|
prng_seed=rng, |
|
num_inference_steps=50, |
|
neg_prompt_ids=negative_prompt_ids, |
|
jit=True, |
|
).images |
|
|
|
output_images = pipe.numpy_to_pil(np.asarray(output.reshape((num_samples,) + output.shape[-3:]))) |
|
return output_images |
|
|
|
with gr.Blocks(theme='kfahn/AnimalPose') as demo: |
|
gr.Markdown( |
|
""" |
|
# Animal Pose Control Net |
|
## This is a demo of Animal Pose ControlNet, which is a model trained on runwayml/stable-diffusion-v1-5 with new type of conditioning. |
|
[Dataset](https://huggingface.co/datasets/JFoz/dog-poses-controlnet-dataset) |
|
[Diffusers model](https://huggingface.co/JFoz/dog-pose) |
|
[Github](https://github.com/fi4cr/animalpose) |
|
[Training Report](https://wandb.ai/john-fozard/AP10K-pose/runs/wn89ezaw) |
|
""") |
|
with gr.Row(): |
|
with gr.Column(): |
|
prompts = gr.Textbox(label="Prompt", value="best quality, extremely detailed") |
|
negative_prompts = gr.Textbox(label="Negative Prompt", value="lowres, bad anatomy, missing ears, missing paws") |
|
conditioning_image = gr.Image(label="Conditioning Image") |
|
run_btn = gr.Button("Run") |
|
with gr.Column(): |
|
|
|
keypoint_tool = gr.HTML(lines) |
|
|
|
run_btn.click(fn=infer, inputs = ["prompts", "negative_prompts", "conditioning_image"], outputs = "gallery") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
demo.launch() |
