Update app_qwen.py

app_qwen.py

@@ -9,19 +9,20 @@ import time
 import os
 
 
-
 @st.cache_resource
 def load_model():
     """Load the model and processor (cached to prevent reloading)"""
+    device = "cuda" if torch.cuda.is_available() else "cpu"
     model = Qwen2VLForConditionalGeneration.from_pretrained(
         "Qwen/Qwen2-VL-2B-Instruct-GPTQ-Int4", 
-        torch_dtype=torch.bfloat16, 
-        device_map="auto"
-    ).eval()
+        torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32, 
+        device_map="auto" if device == "cuda" else None
+    ).eval().to(device)
     processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct-GPTQ-Int4")
-    return model, processor
+    return model, processor, device
+
 
-def process_image(image, prompt, model, processor):
+def process_image(image, prompt, model, processor, device):
     """Process the image and return the model's output"""
     start_time = time.time()
     
@@ -36,7 +37,7 @@ def process_image(image, prompt, model, processor):
     ]
 
     text_prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
-    inputs = processor(text=[text_prompt], images=[image], padding=True, return_tensors="pt").to("cpu")
+    inputs = processor(text=[text_prompt], images=[image], padding=True, return_tensors="pt").to(device)
 
     output_ids = model.generate(**inputs, max_new_tokens=100)
     generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, output_ids)]
@@ -45,13 +46,14 @@ def process_image(image, prompt, model, processor):
     inference_time = time.time() - start_time
     return output_text[0].strip(), inference_time
 
+
 def main():
     # Compact header
     st.markdown("<h1 style='font-size: 24px;'>🔍 Image Analysis with Qwen2-VL</h1>", unsafe_allow_html=True)
 
     # Load model and processor
     with st.spinner("Loading model... This might take a minute."):
-        model, processor = load_model()
+        model, processor, device = load_model()
 
     # Initialize session state
     if 'selected_image' not in st.session_state:
@@ -99,7 +101,7 @@ def main():
     if analyze_button and image_source:
         with st.spinner("Analyzing..."):
             try:
-                result, inference_time = process_image(image, prompt, model, processor)
+                result, inference_time = process_image(image, prompt, model, processor, device)
                 st.session_state.result = result
                 st.session_state.inference_time = inference_time
             except Exception as e:
@@ -144,4 +146,4 @@ def main():
             st.error("No example images found in the 'images' directory")
 
 if __name__ == "__main__":
-    main() 
+    main()