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()