Update app.py

app.py

@@ -1,10 +1,10 @@
 import gradio as gr
 import torch, random, time
 from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
-device = "cuda" if torch.cuda.is_available() else "cpu"
+from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
 translations = {
 	'en': {
-		'model': 'Model Path',
+		'model_name': 'Model Path',
 		'loading': 'Loading',
 		'input': 'Input Image',
 		'prompt': 'Prompt',
@@ -18,7 +18,7 @@ translations = {
 		'seed': 'Seed',
 	},
 	'zh': {
-		'model': '模型路径',
+		'model_name': '模型路径',
 		'loading': '载入',
 		'input': '输入图像',
 		'prompt': '提示',
@@ -37,7 +37,7 @@ def generate_new_seed():
 	return random.randint(1, 2147483647)
 def update_language(new_language):
 	return [
-		gr.Textbox.update(placeholder=translations[new_language]['model']),
+		gr.Textbox.update(placeholder=translations[new_language]['model_name']),
 		gr.Button.update(value=translations[new_language]['loading']),
 		gr.Image.update(label=translations[new_language]['input']),
 		gr.Textbox.update(placeholder=translations[new_language]['prompt']),
@@ -56,7 +56,7 @@ 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))
+	generator = torch.Generator().manual_seed(int(seed))
 	global text2img, img2img
 	start_time = time.time()
 	if image_input is None:
@@ -65,20 +65,25 @@ def Generate(image_input, prompt, negative_prompt, strength, guidance_scale, num
 		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):
+def Loading(model_name, is_xl, is_cuda):
 	global text2img, img2img
-	if device == "cuda":
-		text2img = StableDiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to(device)
+	device = "cuda" if is_xl else "cpu"
+	pipeline_class = StableDiffusionXLPipeline if is_xl else StableDiffusionPipeline
+	if is_xl:
+		text2img = pipeline_class.from_pretrained(model_name, torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to(device)
+	else:
+		text2img = pipeline_class.from_pretrained(model_name, use_safetensors=True).to(device)
+	if is_xl:
 		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
+	img2img = (StableDiffusionXLImg2ImgPipeline if is_xl else StableDiffusionImg2ImgPipeline)(**text2img.components)
+	return model_name
 with gr.Blocks() as demo:
 	with gr.Row():
-		model = gr.Textbox(value="nota-ai/bk-sdm-tiny-2m", label=translations[language]['model'])
+		model_name = gr.Textbox(value="nota-ai/bk-sdm-tiny-2m", label=translations[language]['model_name'])
+		is_xl = gr.Checkbox(label="SDXL")
+		is_cuda = gr.Checkbox(label="cuda", value=torch.cuda.is_available())
 		loading = gr.Button(translations[language]['loading'])
 		set_language = gr.Dropdown(list(translations.keys()), label="Language", value=language)
 	with gr.Row():
@@ -104,6 +109,6 @@ with gr.Blocks() as demo:
 			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])
+	loading.click(Loading, [model_name, is_xl, is_cuda], model_name)
+	set_language.change(update_language, set_language, [model_name, loading, image_input, prompt, negative_prompt, generate, strength, guidance_scale, num_inference_steps, width, height, seed])
 demo.queue().launch()