Update app.py

app.py

@@ -14,6 +14,11 @@ import numpy as np
 from PIL import Image
 import edge_tts
 
+import sambanova_gradio
+# Load the reasoning model from sambanova_gradio.
+# This returns a callable interface for inference.
+reasoning_model = gr.load("DeepSeek-R1-Distill-Llama-70B", src=sambanova_gradio.registry, accept_token=True)
+
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -24,14 +29,6 @@ from transformers import (
 from transformers.image_utils import load_image
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
-# Load the reasoning model interface from sambanova_gradio
-try:
-    import sambanova_gradio
-    reasoning_interface = gr.load("DeepSeek-R1-Distill-Llama-70B", src=sambanova_gradio.registry, accept_token=True)
-except Exception as e:
-    reasoning_interface = None
-    print("Reasoning model could not be loaded:", e)
-
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
@@ -194,8 +191,8 @@ def generate(
 ):
     text = input_dict["text"]
     files = input_dict.get("files", [])
-    lower_text = text.lower().strip()
 
+    lower_text = text.lower().strip()
     # Check if the prompt is an image generation command using model flags.
     if (lower_text.startswith("@lightningv5") or 
         lower_text.startswith("@lightningv4") or 
@@ -248,16 +245,14 @@ def generate(
         yield gr.Image(image_path)
         return
 
-    # New reasoning feature.
+    # New reasoning branch.
     elif lower_text.startswith("@reasoning"):
-        prompt_clean = re.sub(r"@reasoning", "", text, flags=re.IGNORECASE).strip()
-        if reasoning_interface is None:
-            yield "Reasoning model is not available."
-            return
-        yield "Reasoning in progress..."
-        # Call the reasoning model's prediction.
-        result = reasoning_interface.predict(prompt_clean)
-        yield result
+        # Remove the reasoning flag and clean the prompt.
+        prompt_clean = re.sub(r"@reasoning", "", text, flags=re.IGNORECASE).strip().strip('"')
+        yield "Processing reasoning request..."
+        # Call the reasoning model (this call might be synchronous; adjust if needed).
+        reasoning_response = reasoning_model(prompt_clean)
+        yield reasoning_response
         return
 
     # Otherwise, handle text/chat (and TTS) generation.