with to(torch.bfloat16 & logging level

app.py

@@ -1,5 +1,5 @@
 import traceback
-from logging import getLogger
+import logging
 from typing import Optional
 import spaces
 import gradio as gr
@@ -22,7 +22,10 @@ import torch
 from transformers import AutoModelForCausalLM
 from transformers import AutoProcessor 
 
-logger = getLogger(__name__)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+handler = logging.StreamHandler()
+logger.addHandler(handler)
 
 # Define repository and local directory
 repo_id = "microsoft/OmniParser-v2.0"  # HF repo
@@ -36,6 +39,8 @@ magam_model = AutoModelForCausalLM.from_pretrained(magma_model_id, trust_remote_
 magma_processor = AutoProcessor.from_pretrained(magma_model_id, trust_remote_code=True)
 magam_model.to("cuda")
 
+logger.warning(f"The repository is downloading to: {local_dir}")
+
 # Download the entire repository
 snapshot_download(repo_id=repo_id, local_dir=local_dir)
 
@@ -93,11 +98,12 @@ def get_som_response(instruction, image_som):
         add_generation_prompt=True
     )
 
-    inputs = magma_processor(images=[image_som], texts=prompt, return_tensors="pt")
-    inputs['pixel_values'] = inputs['pixel_values'].unsqueeze(0).to(torch.bfloat16) # Add .to(torch.bfloat16) here for explicit casting
-    inputs['image_sizes'] = inputs['image_sizes'].unsqueeze(0)
-    # inputs = inputs.to("cuda")
-    inputs = inputs.to("cuda", dtype=torch.bfloat16)
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        inputs = magma_processor(images=[image_som], texts=prompt, return_tensors="pt")
+        inputs['pixel_values'] = inputs['pixel_values'].unsqueeze(0).to(torch.bfloat16) # Add .to(torch.bfloat16) here for explicit casting
+        inputs['image_sizes'] = inputs['image_sizes'].unsqueeze(0)
+        # inputs = inputs.to("cuda")
+        inputs = inputs.to("cuda", dtype=torch.bfloat16)
 
     magam_model.generation_config.pad_token_id = magma_processor.tokenizer.pad_token_id
     with torch.inference_mode():
@@ -132,11 +138,12 @@ def get_qa_response(instruction, image):
         add_generation_prompt=True
     )
 
-    inputs = magma_processor(images=[image], texts=prompt, return_tensors="pt")
-    inputs['pixel_values'] = inputs['pixel_values'].unsqueeze(0).to(torch.bfloat16) # Add .to(torch.bfloat16) here for explicit casting
-    inputs['image_sizes'] = inputs['image_sizes'].unsqueeze(0)
-    # inputs = inputs.to("cuda")
-    inputs = inputs.to("cuda", dtype=torch.bfloat16)
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        inputs = magma_processor(images=[image], texts=prompt, return_tensors="pt")
+        inputs['pixel_values'] = inputs['pixel_values'].unsqueeze(0).to(torch.bfloat16) # Add .to(torch.bfloat16) here for explicit casting
+        inputs['image_sizes'] = inputs['image_sizes'].unsqueeze(0)
+        # inputs = inputs.to("cuda")
+        inputs = inputs.to("cuda", dtype=torch.bfloat16)
 
     magam_model.generation_config.pad_token_id = magma_processor.tokenizer.pad_token_id
     with torch.inference_mode():