Update app.py

app.py

@@ -11,7 +11,7 @@ transformers.logging.disable_progress_bar()
 warnings.filterwarnings('ignore')
 
 # set device
-torch.set_default_device('cuda:0')  # or 'cpu'
+torch.set_default_device('cuda')  # or 'cpu'
 
 model_name = 'cognitivecomputations/dolphin-vision-7b'
 
@@ -38,14 +38,17 @@ def inference(prompt, image):
     text_chunks = [tokenizer(chunk).input_ids for chunk in text.split('<image>')]
     input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1], dtype=torch.long).unsqueeze(0)
 
-    image_tensor = model.process_images([image], model.config).to(dtype=model.dtype, device=model.device) # Move to model's device
 
-    # generate
-    output_ids = model.generate(
-        input_ids,
-        images=image_tensor,
-        max_new_tokens=2048,
-        use_cache=True)[0]
+    image_tensor = model.process_images([image], model.config).to(dtype=model.dtype, device=model.device)
+
+    # Generate with autocast for mixed precision on GPU
+    with torch.cuda.amp.autocast(): 
+        output_ids = model.generate(
+            input_ids.to(model.device), # Move input_ids to GPU
+            images=image_tensor,
+            max_new_tokens=2048,
+            use_cache=True
+        )[0]
 
     return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()