Update app.py

app.py

@@ -5,13 +5,13 @@ import cv2
 import gradio as gr
 import numpy as np
 from huggingface_hub import snapshot_download
-from transformers import CLIPVisionModelWithProjection,CLIPImageProcessor
+from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor, pipeline
 from diffusers.utils import load_image
 from kolors.pipelines.pipeline_controlnet_xl_kolors_img2img import StableDiffusionXLControlNetImg2ImgPipeline
 from kolors.models.modeling_chatglm import ChatGLMModel
 from kolors.models.tokenization_chatglm import ChatGLMTokenizer
 from kolors.models.controlnet import ControlNetModel
-from diffusers import  AutoencoderKL
+from diffusers import AutoencoderKL
 from kolors.models.unet_2d_condition import UNet2DConditionModel
 from diffusers import EulerDiscreteScheduler
 from PIL import Image
@@ -19,13 +19,15 @@ from annotator.midas import MidasDetector
 from annotator.dwpose import DWposeDetector
 from annotator.util import resize_image, HWC3
 
-
 device = "cuda"
 ckpt_dir = snapshot_download(repo_id="Kwai-Kolors/Kolors")
 ckpt_dir_depth = snapshot_download(repo_id="Kwai-Kolors/Kolors-ControlNet-Depth")
 ckpt_dir_canny = snapshot_download(repo_id="Kwai-Kolors/Kolors-ControlNet-Canny")
 ckpt_dir_pose = snapshot_download(repo_id="Kwai-Kolors/Kolors-ControlNet-Pose")
 
+# Add translation pipeline
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
+
 text_encoder = ChatGLMModel.from_pretrained(f'{ckpt_dir}/text_encoder', torch_dtype=torch.float16).half().to(device)
 tokenizer = ChatGLMTokenizer.from_pretrained(f'{ckpt_dir}/text_encoder')
 vae = AutoencoderKL.from_pretrained(f"{ckpt_dir}/vae", revision=None).half().to(device)
@@ -37,7 +39,7 @@ controlnet_pose = ControlNetModel.from_pretrained(f"{ckpt_dir_pose}", revision=N
 
 pipe_depth = StableDiffusionXLControlNetImg2ImgPipeline(
     vae=vae,
-    controlnet = controlnet_depth,
+    controlnet=controlnet_depth,
     text_encoder=text_encoder,
     tokenizer=tokenizer,
     unet=unet,
@@ -47,7 +49,7 @@ pipe_depth = StableDiffusionXLControlNetImg2ImgPipeline(
 
 pipe_canny = StableDiffusionXLControlNetImg2ImgPipeline(
     vae=vae,
-    controlnet = controlnet_canny,
+    controlnet=controlnet_canny,
     text_encoder=text_encoder,
     tokenizer=tokenizer,
     unet=unet,
@@ -57,7 +59,7 @@ pipe_canny = StableDiffusionXLControlNetImg2ImgPipeline(
 
 pipe_pose = StableDiffusionXLControlNetImg2ImgPipeline(
     vae=vae,
-    controlnet = controlnet_pose,
+    controlnet=controlnet_pose,
     text_encoder=text_encoder,
     tokenizer=tokenizer,
     unet=unet,
@@ -65,6 +67,13 @@ pipe_pose = StableDiffusionXLControlNetImg2ImgPipeline(
     force_zeros_for_empty_prompt=False
 )
 
+@spaces.GPU
+def translate_korean_to_english(text):
+    if any(ord(char) >= 0xAC00 and ord(char) <= 0xD7A3 for char in text):  # Check if Korean characters are present
+        translated = translator(text, max_length=512)[0]['translation_text']
+        return translated
+    return text
+
 @spaces.GPU
 def process_canny_condition(image, canny_threods=[100,200]):
     np_image = image.copy()
@@ -90,7 +99,7 @@ def process_dwpose_condition(image, res=1024):
     img = resize_image(HWC3(image), res)
     out_res, out_img = model_dwpose(image) 
     result = HWC3(out_img)
-    result = cv2.resize( result, (w,h) )
+    result = cv2.resize(result, (w,h))
     return Image.fromarray(result)
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -99,7 +108,7 @@ MAX_IMAGE_SIZE = 1024
 @spaces.GPU
 def infer_depth(prompt, 
           image = None, 
-          negative_prompt = "nsfw，脸部阴影，低分辨率，jpeg伪影、模糊、糟糕，黑脸，霓虹灯", 
+          negative_prompt = "nsfw, facial shadows, low resolution, jpeg artifacts, blurry, bad quality, dark face, neon lights", 
           seed = 397886929, 
           randomize_seed = False,
           guidance_scale = 6.0, 
@@ -108,22 +117,25 @@ def infer_depth(prompt,
           control_guidance_end = 0.9,
           strength = 1.0
         ):
+    prompt = translate_korean_to_english(prompt)
+    negative_prompt = translate_korean_to_english(negative_prompt)
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
-    init_image = resize_image(image,  MAX_IMAGE_SIZE)
+    init_image = resize_image(image, MAX_IMAGE_SIZE)
     pipe = pipe_depth.to("cuda")
-    condi_img = process_depth_condition_midas( np.array(init_image), MAX_IMAGE_SIZE)
+    condi_img = process_depth_condition_midas(np.array(init_image), MAX_IMAGE_SIZE)
     image = pipe(
-        prompt= prompt ,
-        image = init_image,
-        controlnet_conditioning_scale = controlnet_conditioning_scale,
-        control_guidance_end = control_guidance_end, 
-        strength= strength , 
-        control_image = condi_img,
-        negative_prompt= negative_prompt , 
-        num_inference_steps= num_inference_steps, 
-        guidance_scale= guidance_scale,
+        prompt=prompt,
+        image=init_image,
+        controlnet_conditioning_scale=controlnet_conditioning_scale,
+        control_guidance_end=control_guidance_end, 
+        strength=strength, 
+        control_image=condi_img,
+        negative_prompt=negative_prompt, 
+        num_inference_steps=num_inference_steps, 
+        guidance_scale=guidance_scale,
         num_images_per_prompt=1,
         generator=generator,
     ).images[0]
@@ -132,7 +144,7 @@ def infer_depth(prompt,
 @spaces.GPU
 def infer_canny(prompt, 
           image = None, 
-          negative_prompt = "nsfw，脸部阴影，低分辨率，jpeg伪影、模糊、糟糕，黑脸，霓虹灯", 
+          negative_prompt = "nsfw, facial shadows, low resolution, jpeg artifacts, blurry, bad quality, dark face, neon lights", 
           seed = 397886929, 
           randomize_seed = False,
           guidance_scale = 6.0, 
@@ -141,22 +153,25 @@ def infer_canny(prompt,
           control_guidance_end = 0.9,
           strength = 1.0
         ):
+    prompt = translate_korean_to_english(prompt)
+    negative_prompt = translate_korean_to_english(negative_prompt)
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
-    init_image = resize_image(image,  MAX_IMAGE_SIZE)
+    init_image = resize_image(image, MAX_IMAGE_SIZE)
     pipe = pipe_canny.to("cuda")
     condi_img = process_canny_condition(np.array(init_image))
     image = pipe(
-        prompt= prompt ,
-        image = init_image,
-        controlnet_conditioning_scale = controlnet_conditioning_scale,
-        control_guidance_end = control_guidance_end, 
-        strength= strength , 
-        control_image = condi_img,
-        negative_prompt= negative_prompt , 
-        num_inference_steps= num_inference_steps, 
-        guidance_scale= guidance_scale,
+        prompt=prompt,
+        image=init_image,
+        controlnet_conditioning_scale=controlnet_conditioning_scale,
+        control_guidance_end=control_guidance_end, 
+        strength=strength, 
+        control_image=condi_img,
+        negative_prompt=negative_prompt, 
+        num_inference_steps=num_inference_steps, 
+        guidance_scale=guidance_scale,
         num_images_per_prompt=1,
         generator=generator,
     ).images[0]
@@ -165,7 +180,7 @@ def infer_canny(prompt,
 @spaces.GPU
 def infer_pose(prompt, 
           image = None, 
-          negative_prompt = "nsfw，脸部阴影，低分辨率，jpeg伪影、模糊、糟糕，黑脸，霓虹灯", 
+          negative_prompt = "nsfw, facial shadows, low resolution, jpeg artifacts, blurry, bad quality, dark face, neon lights", 
           seed = 66, 
           randomize_seed = False,
           guidance_scale = 6.0, 
@@ -174,45 +189,48 @@ def infer_pose(prompt,
           control_guidance_end = 0.9,
           strength = 1.0
         ):
+    prompt = translate_korean_to_english(prompt)
+    negative_prompt = translate_korean_to_english(negative_prompt)
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
-    init_image = resize_image(image,  MAX_IMAGE_SIZE)
+    init_image = resize_image(image, MAX_IMAGE_SIZE)
     pipe = pipe_pose.to("cuda")
     condi_img = process_dwpose_condition(np.array(init_image), MAX_IMAGE_SIZE)
     image = pipe(
-        prompt= prompt ,
-        image = init_image,
-        controlnet_conditioning_scale = controlnet_conditioning_scale,
-        control_guidance_end = control_guidance_end, 
-        strength= strength , 
-        control_image = condi_img,
-        negative_prompt= negative_prompt , 
-        num_inference_steps= num_inference_steps, 
-        guidance_scale= guidance_scale,
+        prompt=prompt,
+        image=init_image,
+        controlnet_conditioning_scale=controlnet_conditioning_scale,
+        control_guidance_end=control_guidance_end, 
+        strength=strength, 
+        control_image=condi_img,
+        negative_prompt=negative_prompt, 
+        num_inference_steps=num_inference_steps, 
+        guidance_scale=guidance_scale,
         num_images_per_prompt=1,
         generator=generator,
     ).images[0]
     return [condi_img, image], seed
 
 canny_examples = [
-    ["一个漂亮的女孩，高品质，超清晰，色彩鲜艳，超高分辨率，最佳品质，8k，高清，4K",
+    ["아름다운 소녀, 고품질, 초고해상도, 생생한 색상, 최고의 품질, 8k, HD, 4K",
      "image/woman_1.png"],
-    ["全景，一只可爱的白色小狗坐在杯子里，看向镜头，动漫风格，3d渲染，辛烷值渲染",
+    ["전경, 귀여운 흰 강아지가 컵에 앉아 카메라를 보고 있다, 애니메이션 스타일, 3D 렌더링",
     "image/dog.png"]
 ]
 
 depth_examples = [
-    ["新海诚风格，丰富的色彩，穿着绿色衬衫的女人站在田野里，唯美风景，清新明亮，斑驳的光影，最好的质量，超细节，8K画质",
+    ["신카이 마코토 스타일, 풍부한 색감, 녹색 셔츠를 입은 여성이 들판에 서 있다, 아름다운 풍경, 상쾌하고 밝은, 반짝이는 빛, 최고의 품질, 초세밀, 8K 화질",
      "image/woman_2.png"],
-    ["一只颜色鲜艳的小鸟，高品质，超清晰，色彩鲜艳，超高分辨率，最佳品质，8k，高清，4K",
+    ["화려한 색상의 작은 새, 고품질, 초고해상도, 생생한 색상, 최고의 품질, 8k, HD, 4K",
      "image/bird.png"]
 ]
 
 pose_examples = [
-    ["一位穿着紫色泡泡袖连衣裙、戴着皇冠和白色蕾丝手套的女孩双手托脸，高品质，超清晰，色彩鲜艳，超高分辨率，最佳品质，8k，高清，4K",
+    ["보라색 퍼프 소매 드레스를 입고 왕관과 흰색 레이스 장갑을 낀 소녀가 양 손으로 얼굴을 감싸고 있다, 고품질, 초고해상도, 생생한 색상, 최고의 품질, 8k, HD, 4K",
      "image/woman_3.png"],
-    ["一个穿着黑色运动外套、白色内搭，上面戴着项链的女子，站在街边，背景是红色建筑和绿树，高品质，超清晰，色彩鲜艳，超高分辨率，最佳品质，8k，高清，4K",
+    ["검은색 스포츠 재킷과 흰색 이너를 입고 목걸이를 한 여성이 거리에 서 있다, 배경에는 빨간 건물과 녹색 나무가 있다, 고품질, 초고해상도, 생생한 색상, 최고의 품질, 8k, HD, 4K",
      "image/woman_4.png"]
 ]
 
@@ -252,7 +270,7 @@ with gr.Blocks(css=css) as Kolors:
                     label="Negative prompt",
                     placeholder="Enter a negative prompt",
                     visible=True,
-                    value="nsfw，脸部阴影，低分辨率，jpeg伪影、模糊、糟糕，黑脸，霓虹灯"
+                    value="nsfw, facial shadows, low resolution, jpeg artifacts, blurry, bad quality, dark face, neon lights"
                 )
                 seed = gr.Slider(
                     label="Seed",
@@ -353,4 +371,4 @@ with gr.Blocks(css=css) as Kolors:
         outputs = [result, seed_used]
     )
 
-Kolors.queue().launch(debug=True)
+Kolors.queue().launch(debug=True)