Update app.py

app.py

@@ -12,11 +12,15 @@ import spaces
 import torch
 from diffusers import DiffusionPipeline
 from typing import Tuple
+from transformers import pipeline
 
 # Setup rules for bad words (ensure the prompts are kid-friendly)
 bad_words = json.loads(os.getenv('BAD_WORDS', '["violence", "blood", "scary", "death", "ghost"]'))
 default_negative = os.getenv("default_negative","")
 
+# Add the translation pipeline
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
+
 def check_text(prompt, negative=""):
     for i in bad_words:
         if i in prompt:
@@ -120,6 +124,14 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
+def translate_if_korean(text):
+    # Check if the text contains Korean characters
+    if re.search("[\uac00-\ud7a3]", text):
+        # Translate Korean to English
+        translation = translator(text, max_length=512)
+        return translation[0]['translation_text']
+    return text
+
 @spaces.GPU(enable_queue=True)
 def generate(
     prompt: str,
@@ -133,6 +145,9 @@ def generate(
     background: str = "transparent",
     progress=gr.Progress(track_tqdm=True),
 ):
+    # Translate prompt if it's in Korean
+    prompt = translate_if_korean(prompt)
+    
     if check_text(prompt, negative_prompt):
         raise ValueError("Prompt contains restricted words.")
     
@@ -169,18 +184,23 @@ def generate(
     return image_paths, seed
 
 examples = [
-    "cute bunny",
-    "happy cat",
-    "funny dog",
+    "귀여운 고양이",
+    "행복한 토끼",
+    "웃고있는 강아지",
+    "춤추는 돌고래",
+    "신나는 아기 독수리",
+    "즐거운 아기 사자",    
 ]
 
-css = '''
-.gradio-container{max-width: 700px !important}
-h1{text-align:center}
-'''
+css = """
+footer {
+    visibility: hidden;
+}
+"""
+
 
 # Define the Gradio UI for the sticker generator
-with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
+with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:
     gr.Markdown(DESCRIPTION)
     gr.DuplicateButton(
         value="Duplicate Space for private use",