Update app.py

app.py

@@ -1,4 +1,3 @@
-import torch
 from transformers import (
     Qwen2VLForConditionalGeneration, 
     AutoProcessor,
@@ -34,34 +33,26 @@ def load_models():
 
 vision_model, vision_processor, code_model, code_tokenizer = load_models()
 
-VISION_SYSTEM_PROMPT = """You are an AI assistant specialized in analyzing images and videos of code editors. Your primary task is to:
+VISION_SYSTEM_PROMPT = """You are an AI assistant specialized in analyzing images and videos of code editors. Your task is to:
+1. Extract and describe any code snippets visible in the image
+2. Identify any error messages, warnings, or highlighting that indicates bugs
+3. Describe the programming language and context if visible
+Be thorough and accurate in your description, as this will be used to fix the code.
 
-1. FIRST AND MOST IMPORTANTLY: Check if the image contains any inappropriate content such as:
-   - Harassment or bullying
-   - Hate speech or discriminatory content
-   - Sexually explicit material
-   - Dangerous or harmful content
-   If any such content is detected, respond ONLY with: "I apologize, but I cannot process this content as it appears to contain [type of inappropriate content]. Please provide only appropriate code-related images."
+Note: If unwanted contents like Harassment or bullying, Hate speech or discriminatory content, Sexually explicit material, Dangerous or harmful content appears,
+please don't process it and respond ONLY with: "I apologize, but I cannot process this content as it appears to contain [type of inappropriate content]. Please provide only appropriate code-related images."
+"""
 
-2. If the content is appropriate, then:
-   - Extract and describe any code snippets visible in the image
-   - Identify any error messages, warnings, or highlighting that indicates bugs
-   - Describe the programming language and context if visible
+CODE_SYSTEM_PROMPT = """You are an expert code debugging assistant. Based on the description of code and errors provided, your task is to:
+1. Identify the bugs and issues in the code
+2. Provide a corrected version of the code
+3. Explain the fixes made and why they resolve the issues
+4. Provide the output in a well-structured format removing all the unnecessary information
+Be thorough in your explanation and ensure the corrected code is complete and functional.
 
-Be thorough and accurate in your description of appropriate content, as this will be used to fix the code."""
-
-CODE_SYSTEM_PROMPT = """You are an expert code debugging assistant. Your tasks in order are:
-
-1. Check if the input description contains any flags for inappropriate content.
+Note: 1. Check if the input description contains any flags for inappropriate content.
    If it does, respond ONLY with: "I apologize, but I cannot process this request as the original content was flagged as inappropriate."
-
-2. If the content is appropriate, then based on the description of code and errors provided:
-   - Identify the bugs and issues in the code
-   - Provide a corrected version of the code
-   - Explain the fixes made and why they resolve the issues
-   - Provide the output in a well-structured format removing all unnecessary information
-
-Be thorough in your explanation and ensure the corrected code is complete and functional."""
+"""
 
 def process_image_for_code(image):
     # First, process with vision model
@@ -101,10 +92,6 @@ def process_image_for_code(image):
         clean_up_tokenization_spaces=False
     )[0]
 
-    # Check if vision model flagged inappropriate content
-    if "I apologize, but I cannot process this content" in vision_description:
-        return vision_description, "No code analysis provided due to inappropriate content."
-
     # Then, use code model to fix the code
     code_messages = [
         {"role": "system", "content": CODE_SYSTEM_PROMPT},
@@ -156,6 +143,7 @@ def process_video_for_code(video_path, max_frames=16, frame_interval=30):
         
     cap.release()
     
+    # Process the first frame for now (you could extend this to handle multiple frames)
     if frames:
         return process_image_for_code(frames[0])
     else:
@@ -166,16 +154,13 @@ def process_content(content):
     if content is None:
         return "Please upload an image or video file of code with errors.", ""
 
-    try:
-        if content.name.lower().endswith(('.png', '.jpg', '.jpeg')):
-            image = Image.open(content.name)
-            vision_output, code_output = process_image_for_code(image)
-        elif content.name.lower().endswith(('.mp4', '.avi', '.mov')):
-            vision_output, code_output = process_video_for_code(content.name)
-        else:
-            return "Unsupported file type. Please provide an image or video file.", ""
-    except Exception as e:
-        return f"An error occurred while processing the file: {str(e)}", ""
+    if content.name.lower().endswith(('.png', '.jpg', '.jpeg')):
+        image = Image.open(content.name)
+        vision_output, code_output = process_image_for_code(image)
+    elif content.name.lower().endswith(('.mp4', '.avi', '.mov')):
+        vision_output, code_output = process_video_for_code(content.name)
+    else:
+        return "Unsupported file type. Please provide an image or video file.", ""
 
     return vision_output, code_output
 
@@ -188,8 +173,8 @@ iface = gr.Interface(
         gr.Code(label="Fixed Code", language="python")
     ],
     title="Vision Code Debugger",
-    description="Upload an image or video of code with errors for AI analysis and fixes. Note: Only appropriate code-related content will be processed."
+    description="Upload an image or video of code with errors, and the AI will analyze and fix the issues."
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()