updated

app.py

@@ -1,49 +1,33 @@
 import os
-import subprocess
-import sys
 import torch
-import gradio as gr
+import streamlit as st
 from PIL import Image
 from deepseek_vl2.serve.inference import load_model, deepseek_generate, convert_conversation_to_prompts
 from deepseek_vl2.models.conversation import SeparatorStyle
 from deepseek_vl2.serve.app_modules.utils import configure_logger, strip_stop_words, pil_to_base64
-from google.colab import files
 
-# Clone the DeepSeek-VL2 repository if not already present
-repo_dir = "/content/DeepSeek-VL2"
-if not os.path.exists(repo_dir):
-    subprocess.run(["git", "clone", "https://github.com/deepseek-ai/DeepSeek-VL2"], check=True)
-    sys.path.append(repo_dir)  # Add the DeepSeek-VL2 repository to the Python path
-else:
-    print("Repository already cloned.")
-
-# Step 2: Install dependencies if not already installed
-subprocess.run([sys.executable, "-m", "pip", "install", "-r", "/content/DeepSeek-VL2/requirements.txt"])
-
-# Step 3: Verify GPU (Optional)
-print("CUDA Available:", torch.cuda.is_available())
-print("Device:", torch.cuda.current_device() if torch.cuda.is_available() else "CPU")
-print("Device Name:", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "N/A")
-
-# Step 4: Define your model and prediction functions
+# Initialize logger
 logger = configure_logger()
 
+# Global variables for model loading
 MODELS = ["deepseek-ai/deepseek-vl2-tiny"]
 DEPLOY_MODELS = {}
 IMAGE_TOKEN = "<image>"
 
+# Load model function
 def fetch_model(model_name: str, dtype=torch.bfloat16):
     global DEPLOY_MODELS
     if model_name not in DEPLOY_MODELS:
-        print(f"Loading {model_name}...")
+        st.write(f"Loading {model_name}...")
         model_info = load_model(model_name, dtype=dtype)
         tokenizer, model, vl_chat_processor = model_info
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         model = model.to(device)
         DEPLOY_MODELS[model_name] = (tokenizer, model, vl_chat_processor)
-        print(f"Loaded {model_name} on {device}")
+        st.write(f"Loaded {model_name} on {device}")
     return DEPLOY_MODELS[model_name]
 
+# Generate prompt with conversation history
 def generate_prompt_with_history(text, images, history, vl_chat_processor, tokenizer, max_length=2048):
     conversation = vl_chat_processor.new_chat_template()
     if history:
@@ -55,6 +39,7 @@ def generate_prompt_with_history(text, images, history, vl_chat_processor, token
     conversation.append_message(conversation.roles[1], "")
     return conversation
 
+# Convert conversation to Gradio-compatible format
 def to_gradio_chatbot(conv):
     ret = []
     for i, (role, msg) in enumerate(conv.messages[conv.offset:]):
@@ -69,14 +54,12 @@ def to_gradio_chatbot(conv):
             ret[-1][-1] = msg
     return ret
 
+# Prediction function
 def predict(text, images, chatbot, history, model_name="deepseek-ai/deepseek-vl2-tiny"):
-    print("Starting predict function...")
     tokenizer, vl_gpt, vl_chat_processor = fetch_model(model_name)
     if not text:
-        print("Empty text input detected.")
         return chatbot, history, "Empty context."
 
-    print("Processing images...")
     pil_images = [Image.open(img).convert("RGB") for img in images] if images else []
     conversation = generate_prompt_with_history(
         text, pil_images, history, vl_chat_processor, tokenizer
@@ -86,7 +69,6 @@ def predict(text, images, chatbot, history, model_name="deepseek-ai/deepseek-vl2
     gradio_chatbot_output = to_gradio_chatbot(conversation)
 
     full_response = ""
-    print("Generating response...")
     try:
         with torch.no_grad():
             for x in deepseek_generate(
@@ -104,38 +86,31 @@ def predict(text, images, chatbot, history, model_name="deepseek-ai/deepseek-vl2
                 response = strip_stop_words(full_response, stop_words)
                 conversation.update_last_message(response)
                 gradio_chatbot_output[-1][1] = response
-                print(f"Yielding partial response: {response[:50]}...")
                 yield gradio_chatbot_output, conversation.messages, "Generating..."
 
-        print("Generation complete.")
         torch.cuda.empty_cache()
         yield gradio_chatbot_output, conversation.messages, "Success"
     except Exception as e:
-        print(f"Error in generation: {str(e)}")
         yield gradio_chatbot_output, conversation.messages, f"Error: {str(e)}"
 
-# Gradio interface for OCR
-def upload_and_process(image):
-    if image is None:
-        return "Please upload an image.", []
-    prompt = "Extract all text from this image exactly as it appears, ensuring the output is in English only. Preserve spaces, bullets, numbers, and all formatting. Do not translate, generate, or include text in any other language. Stop at the last character of the image text."
-    chatbot = []
-    history = []
-    print("Starting upload_and_process...")
-    for chatbot_output, history_output, status in predict(prompt, [image], chatbot, history):
-        print(f"Status: {status}")
-        if status == "Success":
-            return chatbot_output[-1][1], history_output
-    return "Processing failed.", []
+# Streamlit UI setup
+st.title("DeepSeek-VL2 OCR in Colab")
+st.write("Upload an image and get the extracted text.")
 
-# Launch Gradio app
-with gr.Blocks() as demo:
-    gr.Markdown("### DeepSeek-VL2 OCR in Colab")
-    image_input = gr.Image(type="filepath", label="Upload Image")
-    output_text = gr.Textbox(label="Extracted Text")
-    history_state = gr.State([])
-    submit_btn = gr.Button("Extract Text")
-    submit_btn.click(upload_and_process, inputs=image_input, outputs=[output_text, history_state])
+# Image upload
+image_input = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg"])
 
-demo.launch(share=True, debug=True)  # Added debug=True for more Gradio logs
+# Output text
+output_text = st.text_area("Extracted Text", "")
 
+# Handle the image upload and processing
+if image_input:
+    prompt = "Extract all text from this image exactly as it appears, ensuring the output is in English only."
+    chatbot = []
+    history = []
+    for chatbot_output, history_output, status in predict(prompt, [image_input], chatbot, history):
+        if status == "Success":
+            output_text = chatbot_output[-1][1]
+            st.write("Extracted Text:", output_text)
+        else:
+            st.error(f"Error: {status}")