Custom device map to reduce memory consumption

1920x1080 images now cause the demo to OOM, and we can't downscale more
because the text location features work best at that resolution.

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]