Add specific Emu2 loading: bfloat16 weights

kushal-10 · Apr 30, 2024 · ff86080 · ff86080
1 parent ddf47d5
commit ff86080
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/backends/huggingface_multimodal_api.py b/backends/huggingface_multimodal_api.py
@@ -63,7 +63,11 @@ def load_model(model_spec: backends.ModelSpec):
 
     if hasattr(model_spec, 'trust_remote_code'):
         if model_spec['trust_remote_code']:
-            model = model_type.from_pretrained(hf_model_str, device_map="auto", torch_dtype="auto",
+            if model_spec['model_type'] == "Emu2":
+                model = model_type.from_pretrained(hf_model_str, device_map="auto", torch_dtype="bfloat16",
+                                                   trust_remote_code=model_spec['trust_remote_code'])
+            else:
+                model = model_type.from_pretrained(hf_model_str, device_map="auto", torch_dtype="auto",
                                                trust_remote_code=model_spec['trust_remote_code'])
     else:
         model = model_type.from_pretrained(hf_model_str, device_map="auto", torch_dtype="auto") # Load the model