|
import gradio as gr |
|
import torch, random, time |
|
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline |
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
translations = { |
|
'en': { |
|
'model': 'Model Path', |
|
'loading': 'Loading', |
|
'input': 'Input Image', |
|
'prompt': 'Prompt', |
|
'negative_prompt': 'Negative Prompt', |
|
'generate': 'Generate', |
|
'strength': 'Strength', |
|
'guidance_scale': 'Guidance Scale', |
|
'num_inference_steps': 'Number of Inference Steps', |
|
'width': 'Width', |
|
'height': 'Height', |
|
'seed': 'Seed', |
|
}, |
|
'zh': { |
|
'model': '模型路径', |
|
'loading': '载入', |
|
'input': '输入图像', |
|
'prompt': '提示', |
|
'negative_prompt': '负面提示', |
|
'generate': '生成', |
|
'strength': '强度', |
|
'guidance_scale': '指导尺度', |
|
'num_inference_steps': '推理步数', |
|
'width': '宽度', |
|
'height': '高度', |
|
'seed': '种子', |
|
} |
|
} |
|
language='zh' |
|
def generate_new_seed(): |
|
return random.randint(1, 2147483647) |
|
def update_language(new_language): |
|
return [ |
|
gr.Textbox.update(placeholder=translations[new_language]['model']), |
|
gr.Button.update(value=translations[new_language]['loading']), |
|
gr.Image.update(label=translations[new_language]['input']), |
|
gr.Textbox.update(placeholder=translations[new_language]['prompt']), |
|
gr.Textbox.update(placeholder=translations[new_language]['negative_prompt']), |
|
gr.Button.update(value=translations[new_language]['generate']), |
|
gr.Slider.update(label=translations[new_language]['strength']), |
|
gr.Slider.update(label=translations[new_language]['guidance_scale']), |
|
gr.Slider.update(label=translations[new_language]['num_inference_steps']), |
|
gr.Slider.update(label=translations[new_language]['width']), |
|
gr.Slider.update(label=translations[new_language]['height']), |
|
gr.Number.update(label=translations[new_language]['seed']) |
|
] |
|
|
|
text2img = None |
|
img2img = None |
|
def Generate(image_input, prompt, negative_prompt, strength, guidance_scale, num_inference_steps, width, height, seed): |
|
if seed == -1: |
|
seed = generate_new_seed() |
|
generator = torch.Generator(device).manual_seed(int(seed)) |
|
global text2img, img2img |
|
start_time = time.time() |
|
if image_input is None: |
|
image = text2img(prompt=prompt, negative_prompt=negative_prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, width=width, height=height, num_images_per_prompt=1, generator=generator).images[0] |
|
else: |
|
image = img2img(image=image_input, strength=0.75, prompt=prompt, negative_prompt=negative_prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, width=width, height=height, num_images_per_prompt=1, generator=generator).images[0] |
|
minutes, seconds = divmod(round(time.time() - start_time), 60) |
|
return image, f"{minutes:02d}:{seconds:02d}" |
|
def Loading(model): |
|
global text2img, img2img |
|
if device == "cuda": |
|
text2img = StableDiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to(device) |
|
text2img.enable_xformers_memory_efficient_attention() |
|
text2img.vae.enable_xformers_memory_efficient_attention() |
|
else: |
|
text2img = StableDiffusionPipeline.from_pretrained(model, use_safetensors=True).to(device) |
|
text2img.safety_checker = None |
|
img2img = StableDiffusionImg2ImgPipeline(**text2img.components) |
|
return model |
|
with gr.Blocks() as demo: |
|
with gr.Row(): |
|
model = gr.Textbox(value="nota-ai/bk-sdm-tiny-2m", label=translations[language]['model']) |
|
loading = gr.Button(translations[language]['loading']) |
|
set_language = gr.Dropdown(list(translations.keys()), label="Language", value=language) |
|
with gr.Row(): |
|
with gr.Column(): |
|
with gr.Row(): |
|
image_input = gr.Image(label=translations[language]['input']) |
|
with gr.Column(): |
|
prompt = gr.Textbox("space warrior, beautiful, female, ultrarealistic, soft lighting, 8k", placeholder=translations[language]['prompt'], show_label=False, lines=3) |
|
negative_prompt = gr.Textbox("deformed, distorted, disfigured, poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, mutated hands and fingers, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation,lowres,jpeg artifacts,username,logo,signature,watermark,monochrome,greyscale", placeholder=translations[language]['negative_prompt'], show_label=False, lines=3) |
|
generate = gr.Button(translations[language]['generate']) |
|
with gr.Row(): |
|
with gr.Column(): |
|
strength = gr.Slider(minimum=0, maximum=1, value=0.8, step=0.01, label=translations[language]['strength']) |
|
guidance_scale = gr.Slider(minimum=1, maximum=15, value=7.5, step=0.5, label=translations[language]['guidance_scale']) |
|
num_inference_steps = gr.Slider(minimum=1, maximum=100, value=50, step=1, label=translations[language]['num_inference_steps']) |
|
width = gr.Slider(minimum=512, maximum=2048, value=512, step=8, label=translations[language]['width']) |
|
height = gr.Slider(minimum=512, maximum=2048, value=512, step=8, label=translations[language]['height']) |
|
with gr.Row(): |
|
seed = gr.Number(value=-1, label=translations[language]['seed']) |
|
set_seed = gr.Button("🎲") |
|
with gr.Column(): |
|
image_output = gr.Image() |
|
text_output = gr.Textbox(label="time") |
|
set_seed.click(generate_new_seed, None, seed) |
|
generate.click(Generate, [image_input, prompt, negative_prompt, strength, guidance_scale, num_inference_steps, width, height, seed], [image_output, text_output]) |
|
loading.click(Loading, model, model) |
|
set_language.change(update_language, set_language, [model, loading, image_input, prompt, negative_prompt, generate, strength, guidance_scale, num_inference_steps, width, height, seed]) |
|
demo.queue().launch() |
|
|