Update app.py

app.py

@@ -10,48 +10,50 @@ processor = LlavaProcessor.from_pretrained(model_id)
 model = LlavaForConditionalGeneration.from_pretrained(model_id).to("cpu")
 
 # Initialize inference clients
-client_gemma = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
+client_mistral = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
 
 def llava(inputs):
     """Processes an image and text input using Llava."""
-    image = Image.open(inputs["files"][0]).convert("RGB")
-    prompt = f"<|im_start|>user <image>\n{inputs['text']}<|im_end|>"
-    processed = processor(prompt, image, return_tensors="pt").to("cpu")
-    return processed
+    try:
+        image = Image.open(inputs["files"][0]).convert("RGB")
+        prompt = f"<|im_start|>user <image>\n{inputs['text']}<|im_end|>"
+        processed = processor(prompt, image, return_tensors="pt").to("cpu")
+        return processed
+    except Exception as e:
+        print(f"Error in llava function: {e}")
+        return None
 
 def respond(message, history):
     """Generate a response based on text or image input."""
-    if "files" in message and message["files"]:
-        # Handle image + text input
-        inputs = llava(message)
-        streamer = TextIteratorStreamer(skip_prompt=True, skip_special_tokens=True)
-        thread = Thread(target=model.generate, kwargs=dict(inputs=inputs, max_new_tokens=512, streamer=streamer))
-        thread.start()
-
-        buffer = ""
-        for new_text in streamer:
-            buffer += new_text
-            history[-1][1] = buffer  # Update the latest message in history
-            yield history, history  # Yield both chatbot and history for updating
-
-    else:
-        # Handle text-only input
-        user_message = message["text"]
-        history.append([user_message, None])  # Add user's message with a placeholder response
-
-        # Prepare prompt for the language model
-        prompt = [{"role": "user", "content": msg[0]} for msg in history if msg[0]]
-        response = client_gemma.chat_completion(prompt, max_tokens=200)
-
-        # Extract response and update history
-        bot_message = response["choices"][0]["message"]["content"]
-        history[-1][1] = bot_message  # Update the latest message with bot's response
-        yield history, history  # Yield both chatbot and history for updating
-
-def generate_image(prompt):
-    """Generates an image based on the user prompt."""
-    client = InferenceClient("KingNish/Image-Gen-Pro")
-    return client.predict("Image Generation", None, prompt, api_name="/image_gen_pro")
+    try:
+        if "files" in message and message["files"]:
+            # Handle image + text input
+            inputs = llava(message)
+            if inputs is None:
+                raise ValueError("Failed to process image input")
+            
+            streamer = TextIteratorStreamer(skip_prompt=True, skip_special_tokens=True)
+            thread = Thread(target=model.generate, kwargs=dict(inputs=inputs, max_new_tokens=512, streamer=streamer))
+            thread.start()
+            
+            buffer = ""
+            for new_text in streamer:
+                buffer += new_text
+                history[-1][1] = buffer
+                yield history, history
+        else:
+            # Handle text-only input
+            user_message = message["text"]
+            history.append([user_message, None])
+            prompt = [{"role": "user", "content": msg[0]} for msg in history if msg[0]]
+            response = client_mistral.chat_completion(prompt, max_tokens=200)
+            bot_message = response["choices"][0]["message"]["content"]
+            history[-1][1] = bot_message
+            yield history, history
+    except Exception as e:
+        print(f"Error in respond function: {e}")
+        history[-1][1] = f"An error occurred: {str(e)}"
+        yield history, history
 
 # Set up Gradio interface
 with gr.Blocks() as demo: