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,7 +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)
@@ -58,14 +61,34 @@ 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()
-    
-    return output[0]["generated_text"]
+
+
+    # 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 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;
+    }"""
 
 # Create Gradio interface
 demo = gr.Interface(
@@ -87,8 +110,8 @@ demo = gr.Interface(
         ["علي البطل", "غابة سحرية", "Arabic"],
         ["Lila the Ladybug", "a garden full of flowers", "English"],
         ["ليلى الجنية", "حديقة مليئة بالأزهار", "Arabic"]
-    ],
-    theme="default",
+    ], 
+    css = css_style, 
 )
 
 # Launch the Gradio app