Update app.py

app.py

@@ -299,7 +299,7 @@ EXAMPLES = [
 
 
 @spaces.GPU()
-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):
+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):
     try:
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
@@ -310,18 +310,24 @@ def simple_chat(message: dict, file_obj: io.BytesIO = None, temperature: float =
 
         conversation = []
 
-        if file_obj:
-            choice, contents = mode_load(file_obj)
+        if "files" in message and message["files"]:
+            file_content = message["files"][-1]  # Assuming the last file is the one to be processed
+            choice, contents = mode_load(file_content, message["file_name"])
             if choice == "image":
-                conversation.append({"role": "user", "image": contents, "content": message["text"]})
+                conversation.append({"role": "user", "image": contents, "content": message['text']})
             elif choice == "doc":
-                format_msg = contents + "\n\n\n" + "{} files uploaded.\n".format("1") + message["text"]
+                format_msg = contents + "\n\n\n" + "{} files uploaded.\n".format(len(message["files"])) + message['text']
                 conversation.append({"role": "user", "content": format_msg})
         else:
-            conversation.append({"role": "user", "content": message["text"]})
+            # Handle case where no file is uploaded
+            if not history:
+                # You might want to handle this case depending on your use case
+                # raise gr.Error("Please upload an image first.")
+                contents = None
+                conversation.append({"role": "user", "content": message['text']})
 
         input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
-        
+
         generate_kwargs = dict(
             max_length=max_length,
             do_sample=True,
@@ -341,35 +347,37 @@ def simple_chat(message: dict, file_obj: io.BytesIO = None, temperature: float =
         return PlainTextResponse(f"Error: {str(e)}")
 
 
-@app.post("/chat/")
-async def test_endpoint(text: str, file_url: str = None):
-    if not text:
-        raise HTTPException(status_code=400, detail="Missing 'text' in request body")
 
-    if file_url:
-        file_obj = download_file(file_url)
-        if file_obj is None:
-            raise HTTPException(status_code=400, detail="Failed to download file from URL")
+@app.post("/chat/")
+async def test_endpoint(
+    text: str = Form(...), 
+    file: UploadFile = File(None)
+):
+    # Verificar si se ha subido un archivo
+    if file:
+        # Leer el archivo en memoria
+        file_content = BytesIO(await file.read())
+        file_name = file.filename
+        
+        # Construir el mensaje con el archivo y el texto
+        message = {
+            "text": text,
+            "file_content": file_content,
+            "file_name": file_name
+        }
     else:
-        file_obj = None
-
-    print("Mensaje: ")
-    print(text)
-    print("File: ")
-    print(file_obj)
+        # Si no se sube archivo, solo se incluye el texto
+        message = {
+            "text": text,
+            "file_content": None,
+            "file_name": None
+        }
+    
+    # Llamar a la función `simple_chat` con el mensaje
+    response = simple_chat(message)
     
-    response = simple_chat(text, 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)