Update app.py

app.py

@@ -299,7 +299,7 @@ EXAMPLES = [
 
 
 @spaces.GPU()
-def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096, top_p: float = 1, top_k: int = 10, penalty: float = 1.0):
+def simple_chat(message: dict, file_obj: io.BytesIO = None, temperature: float = 0.8, max_length: int = 4096, top_p: float = 1, top_k: int = 10, penalty: float = 1.0):
     try:
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
@@ -310,9 +310,8 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
 
         conversation = []
 
-        if "file" in message and message["file"]:
-            file_contents = io.BytesIO(message["file"])  # Asegúrate de que sea un objeto BytesIO
-            choice, contents = mode_load(file_contents)
+        if file_obj:
+            choice, contents = mode_load(file_obj)
             if choice == "image":
                 conversation.append({"role": "user", "image": contents, "content": message["text"]})
             elif choice == "doc":
@@ -343,30 +342,30 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
 
 
 @app.post("/chat/")
-async def test_endpoint(
-    text: str = Form(...),
-    file: Optional[UploadFile] = File(None)
-):
-    if not text:
+async def test_endpoint(message: dict):
+    if "text" not in message:
         raise HTTPException(status_code=400, detail="Missing 'text' in request body")
 
-    # Lee el archivo si está presente
-    file_contents = None
-    if file:
-        file_contents = await file.read()
-
-    # Construye el diccionario para `simple_chat`
-    message = {
-        "text": text,
-        "file": file_contents
-    }
-
-    print(message)
+    if "file" in message and message["file"]:
+        file_url = message["file"]
+        file_obj = download_file(file_url)
+        if file_obj is None:
+            raise HTTPException(status_code=400, detail="Failed to download file from URL")
+    else:
+        file_obj = None
 
-    # Llama a `simple_chat` con el diccionario
-    response = simple_chat(message)
+    response = simple_chat(message, file_obj)
     return response
 
+def download_file(url: str) -> io.BytesIO:
+    try:
+        response = requests.get(url)
+        response.raise_for_status()
+        return io.BytesIO(response.content)
+    except Exception as e:
+        print(f"Error downloading file: {str(e)}")
+        return None
+
 
 with gr.Blocks(css=CSS, theme="soft", fill_height=True) as demo:
     gr.HTML(TITLE)