Update app.py

app.py

@@ -9,7 +9,7 @@ from tokenizers import Tokenizer
 from huggingface_hub import hf_hub_download
 import json
 from abc import ABC, abstractmethod
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
@@ -170,11 +170,12 @@ class SAM1Model(keras.Model):
 
         self.embed = keras.layers.Embedding(self.cfg['vocab_size'], self.cfg['d_model'], name="embed_tokens")
 
+        # ✅ FIXED: Was using 'ff_num' — now correctly uses 'ff_dim'
         ff_dim = int(self.cfg['d_model'] * self.cfg['ff_mult'])
         block_args = {
             'd_model': self.cfg['d_model'],
             'n_heads': self.cfg['n_heads'],
-            'ff_dim': ff_num,
+            'ff_dim': ff_dim,          # ✅ Correct variable name
             'dropout': self.cfg['dropout'],
             'max_len': self.cfg['max_len'],
             'rope_theta': self.cfg['rope_theta']
@@ -201,7 +202,7 @@ class SAM1Model(keras.Model):
 
 
 # ==============================================================================
-# Helper: Parameter Counting
+# Helper Functions
 # ==============================================================================
 
 def count_parameters(model):
@@ -399,7 +400,7 @@ async def generate_stream(prompt: str, backend, temperature: float) -> AsyncGene
 
 def chat_fn(message, history, model_choice="SAM-X-1-Large", temperature=0.7):
     backend = available_models[model_choice]
-    prompt = f"User: {message}\nSam:   <think>"
+    prompt = f"User: {message}\nSam:    <think>"
     response = ""
     for chunk in generate_stream(prompt, backend, temperature):
         response += chunk
@@ -442,7 +443,7 @@ async def chat_completions(request: ChatCompletionRequest):
     for msg in request.messages:
         prefix = "User" if msg.role.lower() == "user" else "Sam"
         prompt_parts.append(f"{prefix}: {msg.content}")
-    prompt_parts.append("Sam:   <think>")
+    prompt_parts.append("Sam:    <think>")
     prompt = "\n".join(prompt_parts)
 
     async def event_stream():
@@ -466,11 +467,12 @@ async def chat_completions(request: ChatCompletionRequest):
     else:
         full = ""
         async for token in event_stream():
-            if b"[DONE]" not in token.encode():
+            if "[DONE]" not in token:
                 data = json.loads(token.replace("data: ", "").strip())
                 full += data["choices"][0]["delta"]["content"]
         return {"choices": [{"message": {"content": full}}]}
 
+
 @app.get("/v1/models")
 async def list_models():
     return {
@@ -506,5 +508,5 @@ with gr.Blocks(title="SAM-X-1 Chat", theme=gr.themes.Soft()) as demo:
                 ]
             )
 
-# Mount Gradio app on root
+# Mount Gradio app on root path
 app = gr.mount_gradio_app(app, demo, path="/")