Update app.py

app.py

@@ -1,110 +1,216 @@
+import os
+import random
+
 import gradio as gr
-import torch
-from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler
-from huggingface_hub import hf_hub_download
-import spaces
+from gradio_client import Client, file
+import numpy as np
 from PIL import Image
-import requests
+from typing import Tuple
 from translatepy import Translator
 
-translator = Translator()
-
-# Constants
-base = "stabilityai/stable-diffusion-xl-base-1.0"
-repo = "tianweiy/DMD2"
-checkpoints = {
-    "1-Step" : ["dmd2_sdxl_1step_unet_fp16.bin", 1],
-    "4-Step" : ["dmd2_sdxl_4step_unet_fp16.bin", 4],
-}
-loaded = None
-
-CSS = """
-.gradio-container {
-  max-width: 690px !important;
-}
-footer {
-    visibility: hidden;
-}
+MODEL = os.environ.get("MODEL")
+API_URL = "https://api-inference.huggingface.co/models/tianweiy/DMD2"
+DESCRIPTION = """
+# DMD2 文生图
 """
+translator = Translator()
 
-JS = """function () {
-  gradioURL = window.location.href
-  if (!gradioURL.endsWith('?__theme=dark')) {
-    window.location.replace(gradioURL + '?__theme=dark');
-  }
-}"""
-
-
-
-# Ensure model and scheduler are initialized in GPU-enabled function
-if torch.cuda.is_available():
-    unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
-    pipe = DiffusionPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda")
 
 
-# Function 
-@spaces.GPU()
-def generate_image(prompt, ckpt="4-Step"):
-    global loaded
-    
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+MAX_SEED = np.iinfo(np.int32).max
+
+client = Client(MODEL)
+
+style_list = [
+    {
+        "name": "(无风格)",
+        "prompt": "{prompt}",
+    },
+    {
+        "name": "电影",
+        "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
+    },
+    {
+        "name": "摄影",
+        "prompt": "cinematic photo {prompt} . 35mm photograph, film, bokeh, professional, 4k, highly detailed",
+    },
+    {
+        "name": "动画",
+        "prompt": "anime artwork {prompt} . anime style, key visual, vibrant, studio anime,  highly detailed",
+    },
+    {
+        "name": "漫画",
+        "prompt": "manga style {prompt} . vibrant, high-energy, detailed, iconic, Japanese comic style",
+    },
+    {
+        "name": "数绘",
+        "prompt": "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed",
+    },
+    {
+        "name": "像素",
+        "prompt": "pixel-art {prompt} . low-res, blocky, pixel art style, 8-bit graphics",
+    },
+    {
+        "name": "幻想",
+        "prompt": "ethereal fantasy concept art of  {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
+    },
+    {
+        "name": "朋克",
+        "prompt": "neonpunk style {prompt} . cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional",
+    },
+    {
+        "name": "三维",
+        "prompt": "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting",
+    },
+]   
+styles = {k["name"]: (k["prompt"]) for k in style_list}
+print(styles)
+STYLE_NAMES = list(styles.keys())
+DEFAULT_STYLE_NAME = "(无风格)"
+
+def apply_style(style_name: str, positive: str) -> Tuple[str, str]:
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return p.replace("{prompt}", positive), n
+
+def generate(
+    prompt: str,
+    seed: int = 0,
+    width: int = 1024,
+    height: int = 1024,
+    style: str = DEFAULT_STYLE_NAME,
+    num_images: int = 2,
+    randomize_seed: bool = False,
+    progress=gr.Progress(track_tqdm=True),
+):
     prompt = str(translator.translate(prompt, 'English'))
-
     print(prompt)
-    
-    checkpoint = checkpoints[ckpt][0]
-    num_inference_steps = checkpoints[ckpt][1]
-
-    if loaded != num_inference_steps:
-        pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
-        pipe.unet.load_state_dict(torch.load(hf_hub_download(repo, checkpoint), map_location="cuda"))
-        loaded = num_inference_steps
-
-    if loaded == 1:
-        timesteps=[399]
-    else:
-        timesteps=[999, 749, 499, 249]
-
-    results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0, timesteps=timesteps)
-    return results.images[0]
-
+    seed = int(randomize_seed_fn(seed, randomize_seed))
+    # print(client.view_api())
+    result = client.predict(
+		prompt=prompt,
+		seed=seed,
+		height=height,
+		width=width,
+		num_images=num_images,
+		fast_vae_decode=True,
+		api_name="/inference"
+    )
+    images = result[0]
+    print(images)
+    image_paths = []
+  #  List[Dict(image: filepath, caption: str | None)]
+    for img in images:
+        image_paths.append(img["image"])
+    print(image_paths)
+    return image_paths, seed
 
 examples = [
-    "a cat eating a piece of cheese",
-    "a ROBOT riding a BLUE horse on Mars, photorealistic",
-    "Ironman VS Hulk, ultrarealistic",
-    "a CUTE robot artist painting on an easel",
-    "Astronaut in a jungle, cold color palette, oil pastel, detailed, 8k",
-    "An alien holding sign board contain word 'Flash', futuristic, neonpunk",
-    "Kids going to school, Anime style"
+    "镭射眼的秋田犬",
+    "一只吃起司的猫",
+    "太空中骑马的宇航员",
+    "放学回家的学生们，动画风格",
+    "一个可爱的机器人艺术家在画架上绘画，概念艺术",
+    "一位女士的特写，她戴着透明、棱柱形、精致的复仇女神头饰，摆出应有的姿势，棕色肤色"
 ]
 
-
-# Gradio Interface
-
-with gr.Blocks(css=CSS, js=JS, theme="soft") as demo:
-    gr.HTML("<h1><center>Adobe DMD2🦖</center></h1>")
-    gr.HTML("<p><center><a href='https://huggingface.co/tianweiy/DMD2'>DMD2</a> text-to-image generation</center><br><center>Multi-Languages, 4-step is higher quality & 2X slower</center></p>")
+CSS = '''
+.gradio-container{max-width: 560px !important}
+h1{text-align:center}
+footer {
+    visibility: hidden
+}
+'''
+with gr.Blocks(css=CSS, theme="soft") as demo:
+    gr.Markdown(DESCRIPTION)
     with gr.Group():
         with gr.Row():
-            prompt = gr.Textbox(label='Enter Your Prompt', scale=8)
-            ckpt = gr.Dropdown(label='Steps',choices=['1-Step', '4-Step'], value='4-Step', interactive=True)
-            submit = gr.Button(scale=1, variant='primary')
-    img = gr.Image(label='DMD2 Generated Image')    
+            prompt = gr.Text(
+                label="描述",
+                show_label=False,
+                max_lines=1,
+                placeholder="画什么好呢",
+                container=False,
+                scale=2,
+            )
+            run_button = gr.Button("生成", scale=1)
+        result = gr.Gallery(label="作品", columns=1, preview=True)
+    with gr.Accordion("高级选项", open=False):
+        with gr.Row():
+            num_images = gr.Slider(
+                label="数量",
+                minimum=1,
+                maximum=5,
+                step=1,
+                value=2,
+            )
+            seed = gr.Slider(
+                label="种子",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=1,
+                value=0,
+                visible=True
+            )
+            randomize_seed = gr.Checkbox(label="随机种子", value=True)
+        with gr.Row(visible=True):
+            width = gr.Slider(
+                label="宽",
+                minimum=512,
+                maximum=2048,
+                step=8,
+                value=1024,
+            )
+            height = gr.Slider(
+                label="高",
+                minimum=512,
+                maximum=2048,
+                step=8,
+                value=1024,
+            )
+    with gr.Row(visible=True):
+            style_selection = gr.Radio(
+                show_label=True,
+                container=True,
+                interactive=True,
+                choices=STYLE_NAMES,
+                value=DEFAULT_STYLE_NAME,
+                label="风格化",
+            )
+        
+
     gr.Examples(
         examples=examples,
         inputs=prompt,
-        outputs=img,
-        fn=generate_image,
+        outputs=[result, seed],
+        fn=generate,
         cache_examples="lazy",
     )
-
-    prompt.submit(fn=generate_image,
-                 inputs=[prompt, ckpt],
-                 outputs=img,
-                 )
-    submit.click(fn=generate_image,
-                 inputs=[prompt, ckpt],
-                 outputs=img,
-                 )
     
-demo.queue().launch()
+    gr.on(
+        triggers=[
+            prompt.submit,
+            run_button.click,
+        ],
+        fn=generate,
+        inputs=[
+            prompt,
+            seed,
+            width,
+            height,
+            style_selection,
+            num_images,
+            randomize_seed,
+        ],
+        outputs=[result, seed],
+        api_name="run",
+    )
+
+
+        
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch(show_api=False, debug=False)