-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathgradio_text_to_image.py
151 lines (132 loc) · 5.02 KB
/
gradio_text_to_image.py
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
139
140
141
142
143
144
145
146
147
148
149
150
151
import gradio as gr
import jax
import numpy as np
import jax.numpy as jnp
from flax.training import checkpoints
from diffusers import FlaxControlNetModel, FlaxUNet2DConditionModel, FlaxAutoencoderKL, FlaxDDIMScheduler
from codi.controlnet_flax import FlaxControlNetModel
from codi.pipeline_flax_controlnet import FlaxStableDiffusionControlNetPipeline
from transformers import CLIPTokenizer, FlaxCLIPTextModel
from flax.training.common_utils import shard
from flax.jax_utils import replicate
MODEL_NAME = "CompVis/stable-diffusion-v1-4"
unet, unet_params = FlaxUNet2DConditionModel.from_pretrained(
MODEL_NAME,
subfolder="unet",
revision="flax",
dtype=jnp.float32,
)
vae, vae_params = FlaxAutoencoderKL.from_pretrained(
MODEL_NAME,
subfolder="vae",
revision="flax",
dtype=jnp.float32,
)
text_encoder = FlaxCLIPTextModel.from_pretrained(
MODEL_NAME,
subfolder="text_encoder",
revision="flax",
dtype=jnp.float32,
)
tokenizer = CLIPTokenizer.from_pretrained(
MODEL_NAME,
subfolder="tokenizer",
revision="flax",
dtype=jnp.float32,
)
controlnet = FlaxControlNetModel(
in_channels=unet.config.in_channels,
down_block_types=unet.config.down_block_types,
only_cross_attention=unet.config.only_cross_attention,
block_out_channels=unet.config.block_out_channels,
layers_per_block=unet.config.layers_per_block,
attention_head_dim=unet.config.attention_head_dim,
cross_attention_dim=unet.config.cross_attention_dim,
use_linear_projection=unet.config.use_linear_projection,
flip_sin_to_cos=unet.config.flip_sin_to_cos,
freq_shift=unet.config.freq_shift,
)
scheduler = FlaxDDIMScheduler(
num_train_timesteps=1000,
beta_start=0.00085,
beta_end=0.012,
beta_schedule="scaled_linear",
trained_betas=None,
set_alpha_to_one=True,
steps_offset=0,
)
scheduler_state = scheduler.create_state()
pipeline = FlaxStableDiffusionControlNetPipeline(
vae,
text_encoder,
tokenizer,
unet,
controlnet,
scheduler,
None,
None,
dtype=jnp.float32,
)
controlnet_params = checkpoints.restore_checkpoint("experiments/checkpoint_100000.orbax", target=None)
pipeline_params = {
"vae": vae_params,
"unet": unet_params,
"text_encoder": text_encoder.params,
"scheduler": scheduler_state,
"controlnet": controlnet_params,
}
pipeline_params = replicate(pipeline_params)
def infer(seed, prompt, negative_prompt, steps, cfgr):
rng = jax.random.PRNGKey(int(seed))
num_samples = jax.device_count()
rng = jax.random.split(rng, num_samples)
prompt_ids = pipeline.prepare_text_inputs([prompt] * num_samples)
negative_prompt_ids = pipeline.prepare_text_inputs([negative_prompt] * num_samples)
prompt_ids = shard(prompt_ids)
negative_prompt_ids = shard(negative_prompt_ids)
output = pipeline(
prompt_ids=prompt_ids,
image=None,
params=pipeline_params,
prng_seed=rng,
num_inference_steps=int(steps),
guidance_scale=float(cfgr),
neg_prompt_ids=negative_prompt_ids,
jit=True,
).images
output_images = pipeline.numpy_to_pil(np.asarray(output.reshape((num_samples,) + output.shape[-3:])))
return output_images
with gr.Blocks(theme='gradio/soft') as demo:
gr.Markdown("## CoDi: Conditional Diffusion Distillation for Higher-Fidelity and Faster Image Generation")
gr.Markdown("[\[Paper\]](https://arxiv.org/abs/2310.01407) [\[Project Page\]](https://fast-codi.github.io) [\[Code\]](https://github.com/fast-codi/CoDi)")
with gr.Tab("CoDi on Text-to-Image"):
with gr.Row():
with gr.Column():
gr.Radio(["CompVis/stable-diffusion-v1-4"], value="CompVis/stable-diffusion-v1-4", label="baseline model", info="Chose the undistilled baseline model")
with gr.Column():
gr.Radio(["CoDi/text-to-image-v0-1"], value="CoDi/text-to-image-v0-1", label="distilled codi", info="Chose the distilled conditional model")
with gr.Row():
with gr.Column():
prompt_input = gr.Textbox(label="Prompt")
negative_prompt = gr.Textbox(label="Negative Prompt", value="monochrome, lowres, bad anatomy, worst quality, low quality")
seed = gr.Number(label="Seed", value=0)
output = gr.Gallery(label="Output Images")
with gr.Row():
num_inference_steps = gr.Slider(2, 50, value=4, step=1, label="Steps")
guidance_scale = gr.Slider(2.0, 14.0, value=7.5, step=0.5, label='Guidance Scale')
submit_btn = gr.Button(value = "Submit")
inputs = [
seed,
prompt_input,
negative_prompt,
num_inference_steps,
guidance_scale
]
submit_btn.click(fn=infer, inputs=inputs, outputs=[output])
with gr.Row():
gr.Examples(
examples=["oranges", "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"],
inputs=prompt_input,
fn=infer
)
demo.launch(max_threads=1, share=True)