diff --git a/sites/all/modules/ds/modules/ds_extras/ds_extras.ds_fields_info.inc b/sites/all/modules/ds/modules/ds_extras/ds_extras.ds_fields_info.inc
index 0ee5a92..1783f99 100644
--- a/sites/all/modules/ds/modules/ds_extras/ds_extras.ds_fields_info.inc
+++ b/sites/all/modules/ds/modules/ds_extras/ds_extras.ds_fields_info.inc
@@ -23,8 +23,9 @@ function ds_extras_ds_fields_info($entity_type) {
 
   // Flag support.
   if (variable_get('ds_extras_flag') && module_exists('flag')) {
-    if ($entity_type == 'node') {
-      $flags = flag_get_flags('node');
+    $flags = flag_get_flags($entity_type );
+    if (!empty($flags)) {
+      
       foreach ($flags as $name => $flag) {
         $ui_limit = array();
         if (!empty($flag->types)) {
@@ -32,7 +33,8 @@ function ds_extras_ds_fields_info($entity_type) {
             $ui_limit[] = $type . '|*';
           }
         }
-        $fields['node']['ds_flag_' . $name] = array(
+
+        $fields[$entity_type]['ds_flag_' . $name] = array(
           'title' => t('Flag: ' . $flag->get_label('title')),
           'field_type' => DS_FIELD_TYPE_FUNCTION,
           'function' => 'ds_extras_flags_add_flag_link',