fix-oom (#15)

- Custom device map to reduce memory consumption (7fbd1fa2e4186d15e71f63f19ac7285d798d0816)

app.py

@@ -7,7 +7,15 @@ from transformers import FuyuForCausalLM, FuyuProcessor
 model_id = "adept/fuyu-8b"
 dtype = torch.bfloat16
 
-model = FuyuForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=dtype)
+device_map = {
+    "language_model.model.embed_tokens": "cpu",
+    "language_model.model.layers": 0,
+    "language_model.model.final_layernorm": 0,
+    "language_model.lm_head": "cpu",
+    "vision_embed_tokens": "cpu",
+}
+
+model = FuyuForCausalLM.from_pretrained(model_id, device_map=device_map, torch_dtype=dtype)
 processor = FuyuProcessor.from_pretrained(model_id)
 
 CAPTION_PROMPT = "Generate a coco-style caption.\n"
@@ -36,7 +44,7 @@ def pad_to_size(image, canvas_width=1920, canvas_height=1080):
 
 def predict(image, prompt):
     # image = image.convert('RGB')
-    model_inputs = processor(text=prompt, images=[image]).to(device=model.device)
+    model_inputs = processor(text=prompt, images=[image])
 
     generation_output = model.generate(**model_inputs, max_new_tokens=50)
     prompt_len = model_inputs["input_ids"].shape[-1]
@@ -71,7 +79,7 @@ def localize(image, query):
     padded = resize_to_max(image)
     padded = pad_to_size(padded)
 
-    model_inputs = processor(text=prompt, images=[padded]).to(device=model.device)
+    model_inputs = processor(text=prompt, images=[padded])
     
     outputs = model.generate(**model_inputs, max_new_tokens=40)
     post_processed_bbox_tokens = processor.post_process_box_coordinates(outputs)[0]