Update src/app/response.py

src/app/response.py

@@ -1,67 +1,72 @@
-# Necessary imports
-import sys
-import spaces
-
-# Local imports
-from src.config import (
-    device,
-    model_name,
-    system_prompt,
-    sampling,
-    stream,
-    top_p,
-    top_k,
-    temperature,
-    repetition_penalty,
-    max_new_tokens,
-)
-from src.app.model import load_model_and_tokenizer
-from src.logger import logging
-from src.exception import CustomExceptionHandling
-
-
-# Model and tokenizer
-model, tokenizer = load_model_and_tokenizer(model_name, device)
-
-
-@spaces.GPU(duration=120)
-def describe_image(image: str, question: str) -> str:
-    """
-    Generates an answer to a given question based on the provided image and question.
-
-    Args:
-        - image (str): The path to the image file.
-        - question (str): The question text.
-
-    Returns:
-        str: The generated answer to the question.
-    """
-    try:
-        # Message format for the model
-        msgs = [{"role": "user", "content": [image, question]}]
-
-        # Generate the answer
-        answer = model.chat(
-            image=None,
-            msgs=msgs,
-            tokenizer=tokenizer,
-            sampling=sampling,
-            stream=stream,
-            top_p=top_p,
-            top_k=top_k,
-            temperature=temperature,
-            repetition_penalty=repetition_penalty,
-            max_new_tokens=max_new_tokens,
-            system_prompt=system_prompt,
-        )
-
-        # Log the successful generation of the answer
-        logging.info("Answer generated successfully.")
-
-        # Return the answer
-        return "".join(answer)
-
-    # Handle exceptions that may occur during answer generation
-    except Exception as e:
-        # Custom exception handling
-        raise CustomExceptionHandling(e, sys) from e
+# Necessary imports
+import sys
+import gradio as gr
+import spaces
+
+# Local imports
+from src.config import (
+    device,
+    model_name,
+    system_prompt,
+    sampling,
+    stream,
+    top_p,
+    top_k,
+    temperature,
+    repetition_penalty,
+    max_new_tokens,
+)
+from src.app.model import load_model_and_tokenizer
+from src.logger import logging
+from src.exception import CustomExceptionHandling
+
+
+# Model and tokenizer
+model, tokenizer = load_model_and_tokenizer(model_name, device)
+
+
+@spaces.GPU(duration=120)
+def describe_image(image: str, question: str) -> str:
+    """
+    Generates an answer to a given question based on the provided image and question.
+
+    Args:
+        - image (str): The path to the image file.
+        - question (str): The question text.
+
+    Returns:
+        str: The generated answer to the question.
+    """
+    try:
+        # Check if video or question is None
+        if image is None or question is None:
+            raise gr.Error("Image or question cannot be None.")
+
+        # Message format for the model
+        msgs = [{"role": "user", "content": [image, question]}]
+
+        # Generate the answer
+        answer = model.chat(
+            image=None,
+            msgs=msgs,
+            tokenizer=tokenizer,
+            sampling=sampling,
+            stream=stream,
+            top_p=top_p,
+            top_k=top_k,
+            temperature=temperature,
+            repetition_penalty=repetition_penalty,
+            max_new_tokens=max_new_tokens,
+            system_prompt=system_prompt,
+        )
+
+        # Log the successful generation of the answer
+        logging.info("Answer generated successfully.")
+
+        # Return the answer
+        return "".join(answer)
+
+    # Handle exceptions that may occur during answer generation
+    except Exception as e:
+        # Custom exception handling
+        raise CustomExceptionHandling(e, sys) from e