Image-Phased-Consistency-Model / backup.06302024.app.py
awacke1's picture
Rename app.py to backup.06302024.app.py
97e016b verified
import gradio as gr
import torch
from diffusers import StableDiffusionXLPipeline, StableDiffusionPipeline, LCMScheduler
from diffusers.schedulers import TCDScheduler
import spaces
from PIL import Image
SAFETY_CHECKER = True
checkpoints = {
"2-Step": ["pcm_{}_smallcfg_2step_converted.safetensors", 2, 0.0],
"4-Step": ["pcm_{}_smallcfg_4step_converted.safetensors", 4, 0.0],
"8-Step": ["pcm_{}_smallcfg_8step_converted.safetensors", 8, 0.0],
"16-Step": ["pcm_{}_smallcfg_16step_converted.safetensors", 16, 0.0],
"Normal CFG 4-Step": ["pcm_{}_normalcfg_4step_converted.safetensors", 4, 7.5],
"Normal CFG 8-Step": ["pcm_{}_normalcfg_8step_converted.safetensors", 8, 7.5],
"Normal CFG 16-Step": ["pcm_{}_normalcfg_16step_converted.safetensors", 16, 7.5],
"LCM-Like LoRA": [
"pcm_{}_lcmlike_lora_converted.safetensors",
4,
0.0,
],
}
loaded = None
if torch.cuda.is_available():
pipe_sdxl = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch.float16,
variant="fp16",
).to("cuda")
pipe_sd15 = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16"
).to("cuda")
if SAFETY_CHECKER:
from safety_checker import StableDiffusionSafetyChecker
from transformers import CLIPFeatureExtractor
safety_checker = StableDiffusionSafetyChecker.from_pretrained(
"CompVis/stable-diffusion-safety-checker"
).to("cuda")
feature_extractor = CLIPFeatureExtractor.from_pretrained(
"openai/clip-vit-base-patch32"
)
def check_nsfw_images(
images: list[Image.Image],
) -> tuple[list[Image.Image], list[bool]]:
safety_checker_input = feature_extractor(images, return_tensors="pt").to("cuda")
has_nsfw_concepts = safety_checker(
images=[images], clip_input=safety_checker_input.pixel_values.to("cuda")
)
return images, has_nsfw_concepts
@spaces.GPU(enable_queue=True)
def generate_image(
prompt,
ckpt,
num_inference_steps,
progress=gr.Progress(track_tqdm=True),
mode="sdxl",
):
global loaded
checkpoint = checkpoints[ckpt][0].format(mode)
guidance_scale = checkpoints[ckpt][2]
pipe = pipe_sdxl if mode == "sdxl" else pipe_sd15
if loaded != (ckpt + mode):
pipe.load_lora_weights(
"wangfuyun/PCM_Weights", weight_name=checkpoint, subfolder=mode
)
loaded = ckpt + mode
if ckpt == "LCM-Like LoRA":
pipe.scheduler = LCMScheduler()
else:
pipe.scheduler = TCDScheduler(
num_train_timesteps=1000,
beta_start=0.00085,
beta_end=0.012,
beta_schedule="scaled_linear",
timestep_spacing="trailing",
)
results = pipe(
prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale
)
if SAFETY_CHECKER:
images, has_nsfw_concepts = check_nsfw_images(results.images)
if any(has_nsfw_concepts):
gr.Warning("NSFW content detected.")
return Image.new("RGB", (512, 512))
return images[0]
return results.images[0]
def update_steps(ckpt):
num_inference_steps = checkpoints[ckpt][1]
if ckpt == "LCM-Like LoRA":
return gr.update(interactive=True, value=num_inference_steps)
return gr.update(interactive=False, value=num_inference_steps)
css = """
.gradio-container {
max-width: 60rem !important;
}
"""
with gr.Blocks(css=css) as demo:
gr.Markdown(
"""
# Phased Consistency Model
Phased Consistency Model (PCM) is an image generation technique that addresses the limitations of the Latent Consistency Model (LCM) in high-resolution and text-conditioned image generation.
PCM outperforms LCM across various generation settings and achieves state-of-the-art results in both image and video generation.
[[paper](https://huggingface.co/papers/2405.18407)] [[arXiv](https://arxiv.org/abs/2405.18407)] [[code](https://github.com/G-U-N/Phased-Consistency-Model)] [[project page](https://g-u-n.github.io/projects/pcm)]
"""
)
with gr.Group():
with gr.Row():
prompt = gr.Textbox(label="Prompt", scale=8)
ckpt = gr.Dropdown(
label="Select inference steps",
choices=list(checkpoints.keys()),
value="4-Step",
)
steps = gr.Slider(
label="Number of Inference Steps",
minimum=1,
maximum=20,
step=1,
value=4,
interactive=False,
)
ckpt.change(
fn=update_steps,
inputs=[ckpt],
outputs=[steps],
queue=False,
show_progress=False,
)
submit_sdxl = gr.Button("Run on SDXL", scale=1)
submit_sd15 = gr.Button("Run on SD15", scale=1)
img = gr.Image(label="PCM Image")
gr.Examples(
examples=[
[" astronaut walking on the moon", "4-Step", 4],
[
"Photo of a dramatic cliffside lighthouse in a storm, waves crashing, symbol of guidance and resilience.",
"8-Step",
8,
],
[
"Vincent vangogh style, painting, a boy, clouds in the sky",
"Normal CFG 4-Step",
4,
],
[
"Echoes of a forgotten song drift across the moonlit sea, where a ghost ship sails, its spectral crew bound to an eternal quest for redemption.",
"4-Step",
4,
],
[
"Roger rabbit as a real person, photorealistic, cinematic.",
"16-Step",
16,
],
[
"tanding tall amidst the ruins, a stone golem awakens, vines and flowers sprouting from the crevices in its body.",
"LCM-Like LoRA",
4,
],
],
inputs=[prompt, ckpt, steps],
outputs=[img],
fn=generate_image,
cache_examples="lazy",
)
gr.on(
fn=generate_image,
triggers=[ckpt.change, prompt.submit, submit_sdxl.click],
inputs=[prompt, ckpt, steps],
outputs=[img],
)
gr.on(
fn=lambda *args: generate_image(*args, mode="sd15"),
triggers=[submit_sd15.click],
inputs=[prompt, ckpt, steps],
outputs=[img],
)
demo.queue(api_open=False).launch(show_api=False)