Spaces:
Runtime error
Runtime error
File size: 2,784 Bytes
b2c8596 |
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 |
import jax
import jax.numpy as jnp
from flax import jax_utils
from flax.training.common_utils import shard
from PIL import Image
from argparse import Namespace
import gradio as gr
from diffusers import (
FlaxControlNetModel,
FlaxStableDiffusionControlNetPipeline,
)
args = Namespace(
pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5",
revision="non-ema",
from_pt=True,
controlnet_model_name_or_path="Vincent-luo/controlnet-hands",
controlnet_revision=None,
controlnet_from_pt=False,
)
weight_dtype = jnp.float32
controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
args.controlnet_model_name_or_path,
revision=args.controlnet_revision,
from_pt=args.controlnet_from_pt,
dtype=jnp.float32,
)
pipeline, pipeline_params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
args.pretrained_model_name_or_path,
# tokenizer=tokenizer,
controlnet=controlnet,
safety_checker=None,
dtype=weight_dtype,
revision=args.revision,
from_pt=args.from_pt,
)
pipeline_params["controlnet"] = controlnet_params
pipeline_params = jax_utils.replicate(pipeline_params)
rng = jax.random.PRNGKey(0)
num_samples = jax.device_count()
prng_seed = jax.random.split(rng, jax.device_count())
def infer(prompt, negative_prompt, image):
prompts = num_samples * [prompt]
prompt_ids = pipeline.prepare_text_inputs(prompts)
prompt_ids = shard(prompt_ids)
validation_image = Image.fromarray(image).convert("RGB")
processed_image = pipeline.prepare_image_inputs(num_samples * [validation_image])
processed_image = shard(processed_image)
negative_prompt_ids = pipeline.prepare_text_inputs([negative_prompt] * num_samples)
negative_prompt_ids = shard(negative_prompt_ids)
images = pipeline(
prompt_ids=prompt_ids,
image=processed_image,
params=pipeline_params,
prng_seed=prng_seed,
num_inference_steps=50,
neg_prompt_ids=negative_prompt_ids,
jit=True,
).images
images = images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:])
return images[0]
with gr.Blocks(theme='gradio/soft') as demo:
gr.Markdown("## Stable Diffusion with Hand Control")
gr.Markdown("In this app, you can find different ControlNets with different filters. ")
with gr.Column():
prompt_input = gr.Textbox(label="Prompt")
negative_prompt = gr.Textbox(label="Negative Prompt")
input_image = gr.Image(label="Input Image")
output_image = gr.Image(label="Output Image")
submit_btn = gr.Button(value = "Submit")
inputs = [prompt_input, negative_prompt, input_image]
submit_btn.click(fn=infer, inputs=inputs, outputs=[output_image])
demo.launch() |