Update app.py

app.py

@@ -7,21 +7,22 @@ from indic_transliteration.sanscript import transliterate
 import langdetect
 import re
 
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+# Initialize clients
+text_client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+image_client = InferenceClient("ijohn07/DALLE-4K")
 
 def detect_language_script(text: str) -> tuple[str, str]:
-    """
-    Detect language and script of the input text.
-    Returns (language_code, script_type)
-    """
+    """Detect language and script of the input text.
+    Returns (language_code, script_type)"""
     try:
         # Use confidence threshold to avoid false detections
         lang_detect = langdetect.detect_langs(text)
-        if lang_detect[0].prob > 0.8:  # Only accept high confidence detections
+        if lang_detect[0].prob > 0.8:
+            # Only accept high confidence detections
             lang = lang_detect[0].lang
         else:
             lang = 'en'  # Default to English if unsure
-            
+        
         script = None
         try:
             script = detect_script(text)
@@ -32,10 +33,8 @@ def detect_language_script(text: str) -> tuple[str, str]:
         return 'en', None
 
 def is_romanized_indic(text: str) -> bool:
-    """
-    Check if text appears to be romanized Indic language.
-    More strict pattern matching.
-    """
+    """Check if text appears to be romanized Indic language.
+    More strict pattern matching."""
     # Common Bengali romanized patterns with word boundaries
     bengali_patterns = [
         r'\b(ami|tumi|apni)\b',  # Common pronouns
@@ -50,13 +49,11 @@ def is_romanized_indic(text: str) -> bool:
     return matches >= 2  # Require at least 2 matches to consider it Bengali
 
 def translate_text(text: str, target_lang='en') -> tuple[str, str, bool]:
-    """
-    Translate text to target language, with more conservative translation logic.
-    """
+    """Translate text to target language, with more conservative translation logic."""
     # Skip translation for very short inputs or basic greetings
     if len(text.split()) <= 2 or text.lower() in ['hello', 'hi', 'hey']:
         return text, 'en', False
-        
+    
     original_lang, script = detect_language_script(text)
     is_transliterated = False
     
@@ -69,12 +66,12 @@ def translate_text(text: str, target_lang='en') -> tuple[str, str, bool]:
         except Exception as e:
             print(f"Translation error: {e}")
             return text, 'en', False
-            
+    
     # Check for romanized Indic text only if it's a longer input
     if original_lang == 'en' and len(text.split()) > 2 and is_romanized_indic(text):
         text = romanized_to_bengali(text)
         return translate_text(text, target_lang)  # Recursive call with Bengali script
-            
+    
     return text, 'en', False
 
 def check_custom_responses(message: str) -> str:
@@ -91,12 +88,47 @@ def check_custom_responses(message: str) -> str:
         "who is ur dev": "sk md saad amin",
         "who is ur developer": "sk md saad amin",
     }
-    
     for pattern, response in custom_responses.items():
         if pattern in message_lower:
             return response
     return None
 
+def is_image_request(message: str) -> bool:
+    """Detect if the message is requesting image generation."""
+    image_triggers = [
+        "generate an image",
+        "create an image",
+        "draw",
+        "make a picture",
+        "generate a picture",
+        "create a picture",
+        "generate art",
+        "create art",
+        "make art",
+        "visualize",
+        "show me",
+    ]
+    message_lower = message.lower()
+    return any(trigger in message_lower for trigger in image_triggers)
+
+def generate_image(prompt: str) -> str:
+    """Generate an image using DALLE-4K model."""
+    try:
+        response = image_client.text_to_image(
+            prompt,
+            parameters={
+                "negative_prompt": "blurry, bad quality, nsfw",
+                "num_inference_steps": 30,
+                "guidance_scale": 7.5
+            }
+        )
+        # Save the image and return the path or base64 string
+        # Note: Implementation depends on how you want to handle the image output
+        return response
+    except Exception as e:
+        print(f"Image generation error: {e}")
+        return None
+
 def romanized_to_bengali(text: str) -> str:
     """Convert romanized Bengali text to Bengali script."""
     bengali_mappings = {
@@ -126,11 +158,11 @@ def romanized_to_bengali(text: str) -> str:
     return text_lower
 
 def respond(
-    message, 
-    history: list[tuple[str, str]], 
-    system_message, 
-    max_tokens, 
-    temperature, 
+    message,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
     top_p,
 ):
     # First check for custom responses
@@ -139,9 +171,25 @@ def respond(
         yield custom_response
         return
 
+    # Check if this is an image generation request
+    if is_image_request(message):
+        try:
+            image = generate_image(message)
+            if image:
+                yield f"Here's your generated image based on: {message}"
+                # You'll need to implement the actual image display logic
+                # depending on your Gradio interface requirements
+                return
+            else:
+                yield "Sorry, I couldn't generate the image. Please try again."
+                return
+        except Exception as e:
+            yield f"An error occurred while generating the image: {str(e)}"
+            return
+
     # Handle translation with more conservative approach
     translated_msg, original_lang, was_transliterated = translate_text(message)
-    
+
     # Prepare conversation history - only translate if necessary
     messages = [{"role": "system", "content": system_message}]
     for val in history:
@@ -156,10 +204,10 @@ def respond(
             messages.append({"role": "assistant", "content": val[1]})
     
     messages.append({"role": "user", "content": translated_msg})
-    
+
     # Get response from model
     response = ""
-    for message in client.chat_completion(
+    for message in text_client.chat_completion(
         messages,
         max_tokens=max_tokens,
         stream=True,
@@ -168,18 +216,19 @@ def respond(
     ):
         token = message.choices[0].delta.content
         response += token
-        
-        # Only translate back if the original was definitely non-English
-        if original_lang != 'en' and len(message.split()) > 2:
-            try:
-                translator = GoogleTranslator(source='en', target=original_lang)
-                translated_response = translator.translate(response)
-                yield translated_response
-            except:
-                yield response
-        else:
+
+    # Only translate back if the original was definitely non-English
+    if original_lang != 'en' and len(message.split()) > 2:
+        try:
+            translator = GoogleTranslator(source='en', target=original_lang)
+            translated_response = translator.translate(response)
+            yield translated_response
+        except:
             yield response
+    else:
+        yield response
 
+# Updated Gradio interface to handle images
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[