get @torch .autocast(device_type=cuda, dtype=torch.bfloat16)

app.py

@@ -1,5 +1,5 @@
 import traceback
-import logging
+from logging import getLogger
 from typing import Optional
 import spaces
 import gradio as gr
@@ -22,8 +22,7 @@ import torch
 from transformers import AutoModelForCausalLM
 from transformers import AutoProcessor 
 
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
+logger = getLogger(__name__)
 
 # Define repository and local directory
 repo_id = "microsoft/OmniParser-v2.0"  # HF repo
@@ -40,7 +39,7 @@ magam_model.to("cuda")
 # Download the entire repository
 snapshot_download(repo_id=repo_id, local_dir=local_dir)
 
-print(f"Repository downloaded to: {local_dir}")
+logger.warning(f"Repository downloaded to: {local_dir}")
 
 
 yolo_model = get_yolo_model(model_path='weights/icon_detect/model.pt')
@@ -79,6 +78,7 @@ DEVICE = torch.device('cuda')
 
 @spaces.GPU
 @torch.inference_mode()
+@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
 def get_som_response(instruction, image_som):
     prompt = magma_som_prompt.format(instruction)
     if magam_model.config.mm_use_image_start_end:
@@ -95,10 +95,9 @@ def get_som_response(instruction, image_som):
 
     inputs = magma_processor(images=[image_som], texts=prompt, return_tensors="pt")
     inputs['pixel_values'] = inputs['pixel_values'].unsqueeze(0)
-    # inputs['pixel_values'] = inputs['pixel_values'].to(torch.bfloat16)
     inputs['image_sizes'] = inputs['image_sizes'].unsqueeze(0)
-    inputs = inputs.to("cuda")
-
+    # 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():
@@ -118,6 +117,7 @@ def get_som_response(instruction, image_som):
 
 @spaces.GPU
 @torch.inference_mode()
+@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
 def get_qa_response(instruction, image):
     prompt = magma_qa_prompt.format(instruction)
     if magam_model.config.mm_use_image_start_end:
@@ -135,7 +135,8 @@ def get_qa_response(instruction, image):
     inputs = magma_processor(images=[image], texts=prompt, return_tensors="pt")
     inputs['pixel_values'] = inputs['pixel_values'].unsqueeze(0)
     inputs['image_sizes'] = inputs['image_sizes'].unsqueeze(0)
-    inputs = inputs.to("cuda")
+    # 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():
@@ -165,7 +166,7 @@ def process(
     instruction,
 ) -> Optional[Image.Image]:
 
-    logger.info("Starting processing.")
+    logger.warning("Starting processing.")
     try:
         # image_save_path = 'imgs/saved_image_demo.png'
         # image_input.save(image_save_path)
@@ -184,7 +185,7 @@ def process(
         parsed_content_list = '\n'.join([f'icon {i}: ' + str(v) for i,v in enumerate(parsed_content_list)])
         
         if len(instruction) == 0:
-            print('finish processing')
+            logger.warning('finish processing')
             image = Image.open(io.BytesIO(base64.b64decode(dino_labled_img)))
             return image, str(parsed_content_list)
 
@@ -210,7 +211,7 @@ def process(
             label_coordinates[key] = [val[0] / image_input.size[0], val[1] / image_input.size[1], val[2] / image_input.size[0], val[3] / image_input.size[1]]
         
         magma_response = get_som_response(instruction, image_som)
-        print("magma repsonse: ", magma_response)
+        logger.warning("magma repsonse: ", magma_response)
 
         # map magma_response into the mark id
         mark_id = extract_mark_id(magma_response)
@@ -258,14 +259,14 @@ def process(
             except:
                 image_som = image_input
 
-        logger.info("finish processing")
+        logger.warning("finish processing")
         return image_som, str(parsed_content_list)
     except Exception as e:
         error_message = traceback.format_exc()
         logger.warning(error_message)
         return image_input, error_message
 
-logger.info("Starting App.")
+logger.warning("Starting App.")
 with gr.Blocks() as demo:
     gr.Markdown(MARKDOWN)
     with gr.Row():