Update app.py

app.py

@@ -11,7 +11,7 @@ transformers.logging.disable_progress_bar()
 warnings.filterwarnings('ignore')
 
 # set device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 model_name = 'cognitivecomputations/dolphin-vision-7b'
 
@@ -19,9 +19,9 @@ model_name = 'cognitivecomputations/dolphin-vision-7b'
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype=torch.float16,
-    device_map='auto', # Keep auto device mapping
+    device_map='auto',
     trust_remote_code=True
-).to(device)  # Explicitly move the model to the device
+)
 
 tokenizer = AutoTokenizer.from_pretrained(
     model_name,
@@ -39,12 +39,12 @@ 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).to(device) # Move input_ids to device
+    input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1], dtype=torch.long).unsqueeze(0).to(device)
 
-    image_tensor = model.process_images([image], model.config).to(dtype=model.dtype, device=device) # Move image_tensor to device
+    image_tensor = model.process_images([image], model.config).to(device)
 
     # generate
-    with torch.cuda.amp.autocast(): # Use autocast for mixed precision
+    with torch.cuda.amp.autocast():
         output_ids = model.generate(
             input_ids,
             images=image_tensor,
@@ -65,4 +65,4 @@ with gr.Blocks() as demo:
 
     submit_button.click(fn=inference, inputs=[prompt_input, image_input], outputs=output_text)
 
-demo.launch() 
+demo.launch(share=True)