Update app.py

app.py

@@ -3,15 +3,17 @@ 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):
-    """runwayml/stable-diffusion-v1-5 stabilityai/stable-diffusion-2-1 prompthero/openjourney dreamlike-art/dreamlike-photoreal-2.0 IDEA-CCNL/Taiyi-Stable-Diffusion-1B-Chinese-v0.1"""
+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
-    # pipe = pipe.to("cuda")
+    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 = ""
@@ -50,6 +52,11 @@ def text2image(model, img_style, prompt_input, negative_prompt_input, guidance_s
             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":
@@ -61,6 +68,12 @@ def text2image(model, img_style, prompt_input, negative_prompt_input, guidance_s
             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 = "像素艺术风格"
@@ -133,7 +146,8 @@ def text2image(model, img_style, prompt_input, negative_prompt_input, guidance_s
             negative_prompt=negative_prompt_input,
             guidance_scale=guidance_scale_input,
             num_images_per_prompt=2,
-            num_inference_steps=num_inference_steps_input).images
+            num_inference_steps=num_inference_steps_input,
+            generator=generator).images
 
     # image.show()
     return image[0], image[1]
@@ -143,11 +157,11 @@ style = """h1 {
 }
 """
 
-img_styles = ["Default", "Chinese Painting", "Chinese Landscape Painting", "Oil Painting", 
-            "Landscape Painting", "Ink Painting", "Watercolour Painting", "Sketch",
-            "Ukiyo-e", "Pixel Art", "Abstract", "Dada", "Expressionism", "Futurism", 
-            "Magic Realism", "Minimalism", "Surrealism", "Neo-Impressionism", "Ghibli", 
-            "Naruto", "Jujutsu Kaisen"]
+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",
@@ -160,6 +174,7 @@ with gr.Blocks(title="Tang Poem Translation Vizualization", css=style, theme=the
         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 ⛲")
@@ -169,6 +184,7 @@ with gr.Blocks(title="Tang Poem Translation Vizualization", css=style, theme=the
             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')
 
@@ -178,11 +194,11 @@ with gr.Blocks(title="Tang Poem Translation Vizualization", css=style, theme=the
             out_img2 = gr.Image(type='pil')
 
         submit_btn.click(text2image,
-                inputs=[model_selection, img_style, prompt_input, neg_prompt_input, g_scale, steps_num],
+                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://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>"
+        "<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()