brnm0327's picture
Update app.py
e236679
from diffusers import StableDiffusionPipeline
import torch
import gradio as gr
import opencc
def text2image(model, img_style, prompt_input, negative_prompt_input, guidance_scale_input, num_inference_steps_input, seed):
"""runwayml/stable-diffusion-v1-5 stabilityai/stable-diffusion-2-1 CompVis/stable-diffusion-v1-4 prompthero/openjourney dreamlike-art/dreamlike-photoreal-2.0 IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1"""
if model == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1 (中文)":
model_id = "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1"
else:
model_id = model
pipe = StableDiffusionPipeline.from_pretrained(model_id) # torch_dtype=torch.float16
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = pipe.to(device)
generator = torch.Generator(device).manual_seed(seed)
if img_style == "Default":
style_p = ""
elif img_style == "Chinese Painting":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "中国画风格"
else:
style_p = "Chinese painting style"
elif img_style == "Chinese Landscape Painting":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "中国山水画风格"
else:
style_p = "Chinese landscape painting style"
elif img_style == "Oil Painting":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "油画风格"
else:
style_p = "Oil painting style"
elif img_style == "Landscape Painting":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "风景画风格"
else:
style_p = "Landscape painting style"
elif img_style == "Ink Painting":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p ="水墨画风格"
else:
style_p = "Ink painting style"
elif img_style == "Watercolour Painting":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "水彩风格"
else:
style_p = "Watercolour painting style"
elif img_style == "Sketch":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "素描风格"
else:
style_p = "Sketch style"
elif img_style == "Cartoon":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "卡通风格"
else:
style_p = "Cartoon style"
elif img_style == "Ukiyo-e":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "浮世绘风格"
else:
style_p = "Ukiyo-e style"
elif img_style == "Van Gogh":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "梵高风格"
else:
style_p = "Van Gogh style"
elif img_style == "Vermeer":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "维米尔风格"
else:
style_p = "Vermeer style"
elif img_style == "Pixel Art":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "像素艺术风格"
else:
style_p = "Pixel Art style"
elif img_style == "Abstract":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "抽象风格"
else:
style_p = "Abstract style"
elif img_style == "Dada":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "达达风格"
else:
style_p = "Dada style"
elif img_style == "Expressionism":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "表现主义风格"
else:
style_p = "Expressionism style"
elif img_style == "Futurism":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "未来主义风格"
else:
style_p = "Futurism style"
elif img_style == "Magic Realism":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "魔幻现实主义风格"
else:
style_p = "Magic Realism style"
elif img_style == "Minimalism":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "极简主义风格"
else:
style_p = "Minimalism style"
elif img_style == "Surrealism":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "超现实主义风格"
else:
style_p = "Surrealism style"
elif img_style == "Neo-Impressionism":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "新印象派风格"
else:
style_p = "Neo-Impressionism style"
elif img_style == "Ghibli":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "吉卜力风格"
else:
style_p = "Ghibli style"
elif img_style == "Naruto":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "火影忍者风格"
else:
style_p = "Naruto style"
elif img_style == "Jujutsu Kaisen":
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
style_p = "咒術回戰风格"
else:
style_p = "Jujutsu Kaisen style"
if model == "prompthero/openjourney":
prompt_input = "mdjrny-v4 style, "+prompt_input
if model_id == "IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1":
converter = opencc.OpenCC('t2s.json')
prompt_input = converter.convert(prompt_input)
image = pipe(prompt=prompt_input.replace("\n", " ")+" "+style_p,
negative_prompt=negative_prompt_input,
guidance_scale=guidance_scale_input,
num_images_per_prompt=2,
num_inference_steps=num_inference_steps_input,
generator=generator).images
# image.show()
return image[0], image[1]
style = """h1 {
text-align: center;
}
"""
img_styles = ["Default", "Chinese Painting", "Chinese Landscape Painting", "Oil Painting",
"Landscape Painting", "Ink Painting", "Watercolour Painting", "Sketch", "Cartoon",
"Ukiyo-e", "Van Gogh", "Vermeer","Pixel Art", "Abstract", "Dada", "Expressionism", "Futurism",
"Magic Realism", "Minimalism", "Surrealism", "Neo-Impressionism", "Ghibli",
"Naruto", "Jujutsu Kaisen"]
theme = gr.themes.Default(primary_hue="violet").set(
loader_color="#33A6B8",
slider_color="#00896C",
)
with gr.Blocks(title="Tang Poem Translation Vizualization", css=style, theme=theme) as demo:
gr.Markdown("# Tang Poem Translation 🦚 Vizualization")
with gr.Row():
with gr.Column():
model_selection = gr.inputs.Radio(["runwayml/stable-diffusion-v1-5",
"stabilityai/stable-diffusion-2-1",
"CompVis/stable-diffusion-v1-4",
"prompthero/openjourney",
"dreamlike-art/dreamlike-photoreal-2.0",
"IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1 (中文)"], label="Model ⛲")
img_style = gr.Dropdown(img_styles, label="Image Style 🏛", value="Default", type="value")
prompt_input = gr.Textbox(lines=4, placeholder="Translation Here...", label="Prompt 🏞 (Tang Poem Translation)")
neg_prompt_input = gr.Textbox(lines=1, placeholder="Unwanted Image Specification...", label="Negative Prompt 🙊")
with gr.Accordion("Advanced Options 🪆", open=False):
g_scale = gr.Slider(1, 20, step=0.5, value=7.5, label="Guidance Scale")
steps_num = gr.Slider(10, 100, step=10, value=50, label="Steps")
seed = gr.Slider(1, 9999999, step=1, value=1024, label="Seed")
submit_btn = gr.Button("Submit", variant='primary')
with gr.Tab("First Image Generation 🌄"):
out_img1 = gr.Image(type='pil')
with gr.Tab("Second Image Generation 🌅"):
out_img2 = gr.Image(type='pil')
submit_btn.click(text2image,
inputs=[model_selection, img_style, prompt_input, neg_prompt_input, g_scale, steps_num, seed],
outputs=[out_img1, out_img2])
gr.HTML(
"<center>Powered by <a href='https://runwayml.com/'>Runway 🎢</a>, <a href='https://stability.ai/'>Stability AI 🤖</a>, <a href='https://github.com/CompVis'>CompVis 👁</a>, <a href='https://prompthero.com/'>PromptHero 🦸</a>, <a href='https://dreamlike.art/'>Dreamlike.art 🔮</a>, and <a href='https://www.idea.edu.cn/research/ccnl.html'>IDEA-CCNL 🔆</a></center>"
)
demo.launch()