diff --git a/src/abstract/CollectionWithCategories.ts b/src/abstract/CollectionWithCategories.ts
index dfbb311e4..52ea3f4c5 100644
--- a/src/abstract/CollectionWithCategories.ts
+++ b/src/abstract/CollectionWithCategories.ts
@@ -11,21 +11,25 @@ interface ModelWithCategoryProps {
 const CATEGORY_KEY = 'category';
 
 export abstract class CollectionWithCategories<T extends Model<ModelWithCategoryProps>> extends Collection<T> {
-  abstract get categories(): Categories;
+  abstract getCategories(): Categories;
 
   initCategory(model: T) {
     let category = model.get(CATEGORY_KEY);
+    const isDefined = category instanceof Category;
 
     // Ensure the category exists and it's not already initialized
-    if (category && !(category instanceof Category)) {
+    if (category && !isDefined) {
       if (isString(category)) {
         category = { id: category, label: category };
       } else if (isObject(category) && !category.id) {
         category.id = category.label;
       }
 
-      const catModel = this.categories.add(category);
+      const catModel = this.getCategories().add(category);
       model.set(CATEGORY_KEY, catModel as any, { silent: true });
+      return catModel;
+    } else if (isDefined) {
+      return category as unknown as Category;
     }
   }
 }
diff --git a/src/trait_manager/index.ts b/src/trait_manager/index.ts
index 0b3022e9e..6089b2fe0 100644
--- a/src/trait_manager/index.ts
+++ b/src/trait_manager/index.ts
@@ -1,7 +1,6 @@
 import { debounce } from 'underscore';
 import { Module } from '../abstract';
 import Categories from '../abstract/ModuleCategories';
-import Category from '../abstract/ModuleCategory';
 import { Model } from '../common';
 import Component from '../dom_components/model/Component';
 import EditorModel from '../editor/model/Editor';
@@ -27,8 +26,6 @@ export default class TraitManager extends Module<TraitManagerConfigModule> {
   view?: TraitsView;
 
   TraitsView = TraitsView;
-  Category = Category;
-  Categories = Categories;
   events = TraitsEvents;
   state = new Model<TraitModuleStateProps>({ traits: [] });
   categories = new Categories();
@@ -124,23 +121,24 @@ export default class TraitManager extends Module<TraitManagerConfigModule> {
   }
 
   /**
-   * Get all available categories.
+   * Get trait categories from the currently selected component.
    * @return {Array<Category>}
    */
   getCategories() {
-    return [...this.categories.models];
+    const cmp = this.state.get('component');
+    const categories = cmp?.traits.categories?.models || [];
+    return [...categories];
   }
 
   render() {
     let { view } = this;
-    const { categories, em } = this;
+    const { em } = this;
     view = new TraitsView(
       {
         el: view?.el,
         collection: [],
         editor: em,
         config: this.getConfig(),
-        categories,
       },
       this.getTypes()
     );
diff --git a/src/trait_manager/model/Trait.ts b/src/trait_manager/model/Trait.ts
index 6bafe89b3..2c8ddb688 100644
--- a/src/trait_manager/model/Trait.ts
+++ b/src/trait_manager/model/Trait.ts
@@ -6,6 +6,7 @@ import EditorModel from '../../editor/model/Editor';
 import TraitView from '../view/TraitView';
 import { isDef } from '../../utils/mixins';
 import { CategoryProperties } from '../../abstract/ModuleCategory';
+import Traits from './Traits';
 
 /** @private */
 export interface TraitProperties {
@@ -123,6 +124,10 @@ export default class Trait extends Model<TraitProperties> {
     this.em = em;
   }
 
+  get parent() {
+    return this.collection as unknown as Traits;
+  }
+
   setTarget(target: Component) {
     if (target) {
       const { name, changeProp, value: initValue, getValue } = this.attributes;
diff --git a/src/trait_manager/model/Traits.ts b/src/trait_manager/model/Traits.ts
index ed9b39156..8271a33d8 100644
--- a/src/trait_manager/model/Traits.ts
+++ b/src/trait_manager/model/Traits.ts
@@ -5,11 +5,13 @@ import Component from '../../dom_components/model/Component';
 import EditorModel from '../../editor/model/Editor';
 import Trait, { TraitProperties } from './Trait';
 import TraitFactory from './TraitFactory';
+import Categories from '../../abstract/ModuleCategories';
 
 export default class Traits extends CollectionWithCategories<Trait> {
   em: EditorModel;
   target!: Component;
   tf: TraitFactory;
+  categories = new Categories();
 
   constructor(coll: TraitProperties[], options: { em: EditorModel }) {
     super(coll);
@@ -25,8 +27,8 @@ export default class Traits extends CollectionWithCategories<Trait> {
     return this.em.Traits;
   }
 
-  get categories() {
-    return this.module.categories;
+  getCategories() {
+    return this.categories;
   }
 
   handleReset(coll: TraitProperties[], { previousModels = [] }: { previousModels?: Trait[] } = {}) {
diff --git a/src/trait_manager/view/TraitsView.ts b/src/trait_manager/view/TraitsView.ts
index 6e8062ff6..99cce4452 100644
--- a/src/trait_manager/view/TraitsView.ts
+++ b/src/trait_manager/view/TraitsView.ts
@@ -7,13 +7,13 @@ import EditorModel from '../../editor/model/Editor';
 import Trait from '../model/Trait';
 import { TraitManagerConfigModule } from '../types';
 import TraitView from './TraitView';
+import Traits from '../model/Traits';
 
 interface TraitsViewProps {
   el?: HTMLElement;
   collection: any[];
   editor: EditorModel;
   config: TraitManagerConfigModule;
-  categories: Categories;
 }
 
 const ATTR_CATEGORIES = 'data-categories';
@@ -24,7 +24,6 @@ export default class TraitsView extends DomainViews {
   em: EditorModel;
   pfx: string;
   ppfx: string;
-  categories: Categories;
   renderedCategories = new Map<string, CategoryView>();
   config: TraitManagerConfigModule;
   traitContClass: string;
@@ -33,6 +32,7 @@ export default class TraitsView extends DomainViews {
   traitsEl?: HTMLElement;
   rendered?: boolean;
   itemsView: TraitManager['types'];
+  collection: Traits;
 
   constructor(props: TraitsViewProps, itemsView: TraitManager['types']) {
     super(props);
@@ -46,9 +46,9 @@ export default class TraitsView extends DomainViews {
     this.ppfx = ppfx;
     this.pfx = ppfx + config.stylePrefix || '';
     this.className = `${this.pfx}traits`;
-    this.categories = props.categories || '';
     this.traitContClass = `${ppfx}traits-c`;
     this.catsClass = `${ppfx}trait-categories`;
+    this.collection = new Traits([], { em });
     this.listenTo(em, 'component:toggled', this.updatedCollection);
     this.updatedCollection();
   }
@@ -61,8 +61,7 @@ export default class TraitsView extends DomainViews {
     const { ppfx, em } = this;
     const comp = em.getSelected();
     this.el.className = `${this.traitContClass}s ${ppfx}one-bg ${ppfx}two-color`;
-    // @ts-ignore
-    this.collection = comp ? comp.get('traits') : [];
+    this.collection = comp?.traits || new Traits([], { em });
     this.render();
   }
 
@@ -85,30 +84,20 @@ export default class TraitsView extends DomainViews {
       attributes: model.get('attributes'),
     });
     const rendered = view.render().el;
-    let category = model.get('category');
-
-    // TODO: this part could be consolidated better with BlocksView.ts
-    // Check for categories
-    if (category && this.categories && !(config as any).ignoreCategories) {
-      if (isString(category)) {
-        category = { id: category, label: category };
-      } else if (isObject(category) && !category.id) {
-        category.id = category.label;
-      }
+    const category = model.parent.initCategory(model);
 
-      const catModel = this.categories.add(category);
-      const catId = catModel.getId();
+    if (category) {
+      const catId = category.getId();
       const categories = this.getCategoriesEl();
       let catView = renderedCategories.get(catId);
-      model.set('category', catModel as any, { silent: true });
 
       if (!catView && categories) {
-        catView = new CategoryView({ model: catModel }, config, 'trait').render();
+        catView = new CategoryView({ model: category }, config, 'trait').render();
         renderedCategories.set(catId, catView);
         categories.appendChild(catView.el);
       }
 
-      catView && catView.append(rendered);
+      catView?.append(rendered);
       return;
     }
 
@@ -136,12 +125,12 @@ export default class TraitsView extends DomainViews {
   }
 
   render() {
-    const { ppfx, catsClass, traitContClass } = this;
+    const { el, ppfx, catsClass, traitContClass } = this;
     const frag = document.createDocumentFragment();
     delete this.catsEl;
     delete this.traitsEl;
     this.renderedCategories = new Map();
-    this.el.innerHTML = `
+    el.innerHTML = `
       <div class="${catsClass}" ${ATTR_CATEGORIES}></div>
       <div class="${traitContClass}" ${ATTR_NO_CATEGORIES}></div>
     `;