Update app.py

app.py

@@ -6,6 +6,9 @@ import spaces
 import gc 
 import torch
 
+from gradio_client import Client
+import pyarabic.araby as araby
+
 # Create the necessary directories
 os.makedirs('.gradio/cached_examples/17', exist_ok=True)
 
@@ -42,6 +45,7 @@ def load_model(model_name):
     
     return generator
 
+client = Client("MohamedRashad/arabic-auto-tashkeel")
 @spaces.GPU
 def generate_kids_story(character, setting, language):
     model_name = get_model_name(language)
@@ -57,15 +61,40 @@ def generate_kids_story(character, setting, language):
 
     messages = [{"role": "user", "content": prompt}]
     output = generator(messages)
-
+    
+    text = output[0]["generated_text"]
+    
     # Delete model and associated objects 
     del generator
     # Run garbage collection
     gc.collect ()
     # Empty CUDA cache
     torch.cuda.empty_cache()
+
+    # Remove diacritics from Arabic text then add tashkeel
+    if language == "Arabic":
+        text = client.predict(
+        		input_text=araby.strip_diacritics(text),
+        		api_name="/infer_shakkala")
     
-    return output[0]["generated_text"]
+    return text
+
+css_style = """
+    body {
+        background-image: url('https://cdna.artstation.com/p/assets/images/images/074/776/904/large/pietro-chiovaro-r1-castle-chp.jpg?1712916847');
+        background-size: cover;
+        background-position: center;
+        color: #fff;  /* General text color */
+        font-family: 'Arial', sans-serif;
+    }
+    .title {
+        text-align: center;
+        font-size: 3.5em;
+        margin: 20px 0;
+        text-shadow: 2px 2px 5px rgba(0, 0, 0, 0.7);
+        color: #FFD700; /* Gold color for title */
+    }
+"""
 
 # Create Gradio interface
 demo = gr.Interface(
@@ -88,7 +117,7 @@ demo = gr.Interface(
         ["Lila the Ladybug", "a garden full of flowers", "English"],
         ["ليلى الجنية", "حديقة مليئة بالأزهار", "Arabic"]
     ],
-    theme="default",
+    css = css_style
 )
 
 # Launch the Gradio app