|
import os |
|
import re |
|
import random |
|
import json |
|
import requests |
|
import gradio as gr |
|
from gradio_client import Client, file |
|
import numpy as np |
|
from PIL import Image |
|
from typing import Tuple |
|
from translatepy import Translator |
|
from io import BytesIO |
|
|
|
MODEL = os.environ.get("MODEL") |
|
MODEL2 = os.environ.get("MODEL2") |
|
MODEL_TEST = os.environ.get("MODEL_TEST") |
|
|
|
DESCRIPTION = """ |
|
# IMGEN🧚♀️ |
|
<center>🧚♂️输入描述生成图片，选择喜欢的风格~<br>🧚默认使用Turbo快速生成，高级选项可更改。</center> |
|
""" |
|
MODEL_ID="stabilityai/sdxl-turbo" |
|
|
|
LICENSE = "采用 SDXL 模型" |
|
translator = Translator() |
|
|
|
|
|
|
|
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) |
|
client2 = Client(MODEL2) |
|
client_test=Client(MODEL_TEST) |
|
|
|
style_list = [ |
|
{ |
|
"name": "(无风格)", |
|
"prompt": "{prompt}", |
|
}, |
|
{ |
|
"name": "电影", |
|
"prompt": "cinematic, {prompt}, detailed, bokeh, moody, film grain", |
|
}, |
|
{ |
|
"name": "摄影", |
|
"prompt": "realistic photo, {prompt}, 35mm photograph, 4k", |
|
}, |
|
{ |
|
"name": "动画", |
|
"prompt": "anime style, {prompt}, studio anime", |
|
}, |
|
{ |
|
"name": "漫画", |
|
"prompt": "manga style, {prompt}, detailed, Japanese comic style", |
|
}, |
|
{ |
|
"name": "数绘", |
|
"prompt": "{prompt}, digital artwork, detailed", |
|
}, |
|
{ |
|
"name": "像素", |
|
"prompt": "pixel-art, {prompt}, 8-bit graphics", |
|
}, |
|
{ |
|
"name": "幻想", |
|
"prompt": "fantasy style, {prompt}, epic, magical, dreamy", |
|
}, |
|
{ |
|
"name": "朋克", |
|
"prompt": "cyberpunk style, {prompt}, stunningly beautiful, ultra detailed, realistic", |
|
}, |
|
{ |
|
"name": "三维", |
|
"prompt": "3d model, {prompt}, highly quality, dramatic lighting", |
|
}, |
|
] |
|
styles = {k["name"]: (k["prompt"]) for k in style_list} |
|
STYLE_NAMES = list(styles.keys()) |
|
DEFAULT_STYLE_NAME = "(无风格)" |
|
|
|
def apply_style(style_name: str, positive: str) -> Tuple[str]: |
|
p = styles.get(style_name, styles[DEFAULT_STYLE_NAME]) |
|
return p.replace("{prompt}", positive) |
|
|
|
def list_dir(path): |
|
for root, dirs, files in os.walk(path): |
|
for file in files: |
|
print(os.path.join(root, file)) |
|
|
|
def generate( |
|
prompt: str, |
|
seed: float = 0, |
|
steps: float = 1, |
|
style: str = DEFAULT_STYLE_NAME, |
|
randomize_seed: bool = False, |
|
use_turbo: bool = True, |
|
use_flash: bool = False, |
|
progress=gr.Progress(track_tqdm=True), |
|
): |
|
prompt = str(translator.translate(prompt, 'English')) |
|
print(f'prompt:{prompt}') |
|
seed = int(randomize_seed_fn(seed, randomize_seed)) |
|
|
|
|
|
prompt = apply_style(style, prompt) |
|
strength = 0.7 |
|
images = [] |
|
if use_turbo and use_flash is False: |
|
result = client.predict( |
|
None, |
|
prompt, |
|
strength, |
|
steps, |
|
seed, |
|
fn_index=1 |
|
) |
|
print(result) |
|
image_url = MODEL+"file="+result |
|
images_request = requests.get(image_url) |
|
print(f'A:{images_request}') |
|
pil_image = Image.open(BytesIO(images_request.content)) |
|
images.append(pil_image) |
|
print(f'B:{images}') |
|
elif use_flash is True: |
|
result = client_test.predict( |
|
prompt, |
|
0, |
|
True, |
|
api_name="/infer" |
|
) |
|
print(result) |
|
image_url = MODEL_TEST+"file="+result |
|
images_request = requests.get(image_url) |
|
print(f'A:{images_request}') |
|
pil_image = Image.open(BytesIO(images_request.content)) |
|
images.append(pil_image) |
|
print(f'B:{images}') |
|
else: |
|
seed = 0 |
|
job = client2.submit( |
|
prompt, |
|
"low_quality", |
|
7.5, |
|
"(No style)", |
|
fn_index=3 |
|
) |
|
result = job.result() |
|
|
|
for root, dirs, files in os.walk(result): |
|
for file in files: |
|
print(file) |
|
if file.endswith('.jpg'): |
|
images.append(os.path.join(root,file)) |
|
|
|
return images, seed |
|
|
|
examples = [ |
|
"坐在火箭上的秋田犬", |
|
"一只吃芝士片的猫", |
|
"太空中骑马的宇航员", |
|
"放学回家的孩子们", |
|
"一个可爱的机器人在照镜子", |
|
"钢铁侠在恐龙时代" |
|
] |
|
|
|
CSS = ''' |
|
.gradio-container{max-width: 560px !important} |
|
h1{text-align:center} |
|
footer { |
|
visibility: hidden |
|
} |
|
''' |
|
with gr.Blocks(css=CSS, theme="xiaobaiyuan/theme_brief") as demo: |
|
gr.Markdown(DESCRIPTION) |
|
with gr.Group(): |
|
with gr.Row(): |
|
prompt = gr.Text( |
|
label="描述", |
|
show_label=False, |
|
max_lines=1, |
|
placeholder="画什么好呢", |
|
container=False, |
|
scale=3, |
|
) |
|
run_button = gr.Button("生成", scale=1) |
|
result = gr.Gallery(label="作品", columns = 1, preview=True) |
|
with gr.Accordion("高级选项", open=False): |
|
with gr.Row(): |
|
steps = gr.Slider( |
|
label="步数", |
|
minimum=1, |
|
maximum=10, |
|
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) |
|
use_turbo = gr.Checkbox(label="快速生成", value=True) |
|
use_flash = gr.Checkbox(label="Flash", value=False) |
|
with gr.Row(visible=True): |
|
style_selection = gr.Radio( |
|
show_label=True, |
|
container=True, |
|
interactive=True, |
|
choices=STYLE_NAMES, |
|
value=DEFAULT_STYLE_NAME, |
|
label="风格化", |
|
) |
|
with gr.Row(): |
|
gr.Examples( |
|
examples=examples, |
|
inputs=prompt, |
|
outputs=[result, seed], |
|
fn=generate, |
|
cache_examples='lazy', |
|
label="示例", |
|
) |
|
with gr.Row(): |
|
gr.Markdown(LICENSE) |
|
|
|
gr.on( |
|
triggers=[ |
|
prompt.submit, |
|
run_button.click, |
|
], |
|
fn=generate, |
|
inputs=[ |
|
prompt, |
|
seed, |
|
steps, |
|
style_selection, |
|
randomize_seed, |
|
use_turbo, |
|
use_flash, |
|
], |
|
outputs=[result, seed], |
|
api_name="run", |
|
) |
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
demo.queue(max_size=20).launch(show_api=False, debug=False) |
