Artiprocher's picture
3x3
473b807
raw
history blame
1.88 kB
import gradio as gr
from LdmZhPipeline import LDMZhTextToImagePipeline
import torch
import numpy as np
from PIL import Image
device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "alibaba-pai/pai-diffusion-poem-large-zh"
pipe_text2img = LDMZhTextToImagePipeline.from_pretrained(model_id, use_auth_token="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")
pipe_text2img = pipe_text2img.to(device)
def infer_text2img(prompt, guide, steps):
output = pipe_text2img([prompt]*9, guidance_scale=guide, num_inference_steps=steps)
images = output.images
images = [np.array(images[i]) for i in range(9)]
images = np.concatenate([
np.concatenate(images[0:3], axis=0),
np.concatenate(images[3:6], axis=0),
np.concatenate(images[6:9], axis=0),
], axis=1)
images = Image.fromarray(images)
return images
with gr.Blocks() as demo:
examples = [
["远上寒山石径斜 白云深处有人家"],
["停车坐爱枫林晚 霜叶红于二月花"],
["接天莲叶无穷碧 映日荷花别样红"],
]
with gr.Row():
with gr.Column(scale=1, ):
image_out = gr.Image(label = '输出(output)')
with gr.Column(scale=1, ):
prompt = gr.Textbox(label = '提示词(prompt)')
submit_btn = gr.Button("生成图像(Generate)")
with gr.Row(scale=0.5 ):
guide = gr.Slider(2, 15, value = 7, label = '文本引导强度(guidance scale)')
steps = gr.Slider(10, 50, value = 20, step = 1, label = '迭代次数(inference steps)')
ex = gr.Examples(examples, fn=infer_text2img, inputs=[prompt, guide, steps], outputs=image_out)
submit_btn.click(fn = infer_text2img, inputs = [prompt, guide, steps], outputs = image_out)
demo.queue(concurrency_count=1, max_size=8).launch()