Update app.py

app.py

@@ -1,39 +1,54 @@
 from __future__ import annotations
 
-import os
+import spaces
+
+import gradio as gr
+from threading import Thread
+from transformers import TextIteratorStreamer
 import hashlib
+import os
+
+from transformers import AutoModel, AutoProcessor
 import torch
-from threading import Thread
-from transformers import AutoModel, AutoProcessor, TextIteratorStreamer
-import gradio as gr
 
-# Initialize the model and processor
-def initialize_model_and_processor():
-    model = AutoModel.from_pretrained("OEvortex/HelpingAI-Vision", torch_dtype=torch.float16, trust_remote_code=True).to("cuda" if torch.cuda.is_available() else "cpu")
-    processor = AutoProcessor.from_pretrained("OEvortex/HelpingAI-Vision", trust_remote_code=True)
-    return model, processor
+model = AutoModel.from_pretrained("OEvortex/HelpingAI-Vision", torch_dtype=torch.float16, trust_remote_code=True).to("cuda")
+
+processor = AutoProcessor.from_pretrained("OEvortex/HelpingAI-Vision", trust_remote_code=True)
+
+if torch.cuda.is_available():
+    DEVICE = "cuda"
+    DTYPE = torch.float16
+else:
+    DEVICE = "cpu"
+    DTYPE = torch.float32
 
-# Function to process images and cache results
 def cached_vision_process(image, max_crops, num_tokens):
     image_hash = hashlib.sha256(image.tobytes()).hexdigest()
     cache_path = f"visual_cache/{image_hash}-{max_crops}-{num_tokens}.pt"
     if os.path.exists(cache_path):
-        return torch.load(cache_path).to(model.device, dtype=model.dtype)
+        return torch.load(cache_path).to(DEVICE, dtype=DTYPE)
     else:
         processor_outputs = processor.image_processor([image], max_crops)
-        pixel_values = [value.to(model.device, model.dtype) for value in processor_outputs["pixel_values"]]
-        coords = [value.to(model.device, model.dtype) for value in processor_outputs["coords"]]
+        pixel_values = processor_outputs["pixel_values"]
+        pixel_values = [
+            value.to(model.device).to(model.dtype) for value in pixel_values
+        ]
+        coords = processor_outputs["coords"]
+        coords = [value.to(model.device).to(model.dtype) for value in coords]
         image_outputs = model.vision_model(pixel_values, coords, num_tokens)
         image_features = model.multi_modal_projector(image_outputs)
         os.makedirs("visual_cache", exist_ok=True)
         torch.save(image_features, cache_path)
-        return image_features.to(model.device, model.dtype)
+        return image_features.to(DEVICE, dtype=DTYPE)
 
-# Function to answer questions about images
+@spaces.GPU(duration=20)
 def answer_question(image, question, max_crops, num_tokens, sample, temperature, top_k):
-    if not question.strip() or not image:
-        return "Please provide both an image and a question."
-    
+    if question is None or question.strip() == "":
+        yield "Please ask me anything"
+        return
+    if image is None:
+        yield "Please upload a picture"
+        return
     prompt = f"""user
 <image>
 {question}
@@ -42,7 +57,6 @@ assistant
     streamer = TextIteratorStreamer(processor.tokenizer, skip_special_tokens=True)
     with torch.inference_mode():
         inputs = processor(prompt, [image], model, max_crops=max_crops, num_tokens=num_tokens)
-    
     generation_kwargs = {
         "input_ids": inputs["input_ids"],
         "attention_mask": inputs["attention_mask"],
@@ -71,15 +85,17 @@ assistant
             yield buffer
     return buffer
 
-# Initialize the model and processor
-model, processor = initialize_model_and_processor()
 
-# Gradio interface setup
 with gr.Blocks() as demo:
     with gr.Group():
         with gr.Row():
-            prompt = gr.Textbox(label="Question", placeholder="e.g. Describe this?", scale=4)
-            submit = gr.Button("Send", scale=1)
+            prompt = gr.Textbox(
+                label="Question", placeholder="e.g. Discribe this?", scale=4
+            )
+            submit = gr.Button(
+                "Send",
+                scale=1,
+            )
         with gr.Row():
             max_crops = gr.Slider(minimum=0, maximum=200, step=5, value=0, label="Max crops")
             num_tokens = gr.Slider(minimum=728, maximum=2184, step=10, value=728, label="Number of image tokens")
@@ -94,4 +110,4 @@ with gr.Blocks() as demo:
     submit.click(answer_question, [img, prompt, max_crops, num_tokens, sample, temperature, top_k], output)
     prompt.submit(answer_question, [img, prompt, max_crops, num_tokens, sample, temperature, top_k], output)
 
-demo.queue().launch(debug=True)
+demo.queue().launch(debug=True)