diff --git a/src/main/java/nl/andrewlalis/erme/control/actions/ExportToImageAction.java b/src/main/java/nl/andrewlalis/erme/control/actions/ExportToImageAction.java
index 7ddeb19..a040a3d 100644
--- a/src/main/java/nl/andrewlalis/erme/control/actions/ExportToImageAction.java
+++ b/src/main/java/nl/andrewlalis/erme/control/actions/ExportToImageAction.java
@@ -114,13 +114,13 @@ public class ExportToImageAction extends AbstractAction {
 		DiagramPanel.prepareGraphics(g2d);
 
 		// Render the model.
-		boolean lolcat = AttributeViewModel.getLolcatMode(); // save previous lolcat mode
-		AttributeViewModel.setLolcatMode(false);
+		boolean lolcat = LolcatAction.getInstance().isLolcatEnabled(); // save previous lolcat mode
+		LolcatAction.getInstance().setLolcatEnabled(false);
 		List<Relation> selectedRelations = this.model.getSelectedRelations();
 		this.model.getSelectedRelations().forEach(r -> r.setSelected(false));
 		new MappingModelViewModel(this.model).draw(g2d);
 		this.model.getRelations().forEach(r -> r.setSelected(selectedRelations.contains(r)));
-		AttributeViewModel.setLolcatMode(lolcat); // revert previous lolcat mode
+		LolcatAction.getInstance().setLolcatEnabled(lolcat); // revert previous lolcat mode
 
 		// Revert back to the normal image space, and render a watermark.
 		g2d.setTransform(originalTransform);
diff --git a/src/main/java/nl/andrewlalis/erme/control/actions/LolcatAction.java b/src/main/java/nl/andrewlalis/erme/control/actions/LolcatAction.java
index d3a9925..b193c96 100644
--- a/src/main/java/nl/andrewlalis/erme/control/actions/LolcatAction.java
+++ b/src/main/java/nl/andrewlalis/erme/control/actions/LolcatAction.java
@@ -1,6 +1,9 @@
 package nl.andrewlalis.erme.control.actions;
 
+import lombok.Getter;
+import lombok.Setter;
 import nl.andrewlalis.erme.EntityRelationMappingEditor;
+import nl.andrewlalis.erme.view.DiagramPanel;
 import nl.andrewlalis.erme.view.view_models.AttributeViewModel;
 
 import javax.swing.*;
@@ -8,6 +11,11 @@ import java.awt.event.ActionEvent;
 
 public class LolcatAction extends AbstractAction {
 	private static LolcatAction instance;
+
+	@Getter
+	@Setter
+	private boolean lolcatEnabled = false;
+
 	public static LolcatAction getInstance() {
 		if (instance == null) {
 			instance = new LolcatAction();
@@ -15,13 +23,16 @@ public class LolcatAction extends AbstractAction {
 		return instance;
 	}
 
+	@Setter
+	private DiagramPanel diagramPanel;
+
 	public LolcatAction() {
 		super("Toggle Lolcat Mode");
 	}
 
 	@Override
 	public void actionPerformed(ActionEvent actionEvent) {
-		AttributeViewModel.setLolcatMode(((AbstractButton)actionEvent.getSource()).getModel().isSelected());
-		EntityRelationMappingEditor.getFrame().getContentPane().repaint();
+		lolcatEnabled = ((AbstractButton)actionEvent.getSource()).getModel().isSelected();
+		diagramPanel.repaint();
 	}
 }
diff --git a/src/main/java/nl/andrewlalis/erme/view/DiagramPanel.java b/src/main/java/nl/andrewlalis/erme/view/DiagramPanel.java
index 8210b1e..54bc65e 100644
--- a/src/main/java/nl/andrewlalis/erme/view/DiagramPanel.java
+++ b/src/main/java/nl/andrewlalis/erme/view/DiagramPanel.java
@@ -128,6 +128,7 @@ public class DiagramPanel extends JPanel implements ModelChangeListener {
 		AddAttributeAction.getInstance().setModel(this.model);
 		RemoveAttributeAction.getInstance().setModel(this.model);
 		LoadSampleModelAction.getInstance().setDiagramPanel(this);
+		LolcatAction.getInstance().setDiagramPanel(this);
 	}
 
 	public static void prepareGraphics(Graphics2D g) {
diff --git a/src/main/java/nl/andrewlalis/erme/view/view_models/AttributeViewModel.java b/src/main/java/nl/andrewlalis/erme/view/view_models/AttributeViewModel.java
index 595b866..6b32638 100644
--- a/src/main/java/nl/andrewlalis/erme/view/view_models/AttributeViewModel.java
+++ b/src/main/java/nl/andrewlalis/erme/view/view_models/AttributeViewModel.java
@@ -1,6 +1,7 @@
 package nl.andrewlalis.erme.view.view_models;
 
 import nl.andrewlalis.erme.EntityRelationMappingEditor;
+import nl.andrewlalis.erme.control.actions.LolcatAction;
 import nl.andrewlalis.erme.model.Attribute;
 import nl.andrewlalis.erme.model.AttributeType;
 import nl.andrewlalis.erme.model.ForeignKeyAttribute;
@@ -22,21 +23,12 @@ public class AttributeViewModel implements ViewModel {
 	public static final float FK_FONT_SIZE = 11.0f;
 	private static final float LOLCAT_SAT = 0.75f;
 	private static final float LOLCAT_BRIGHT = 1f;
-	private static boolean lolcatMode;
 	private final Attribute attribute;
 
 	public AttributeViewModel(Attribute attribute) {
 		this.attribute = attribute;
 	}
 
-	public static boolean getLolcatMode() {
-		return lolcatMode;
-	}
-
-	public static void setLolcatMode(boolean lolcatMode) {
-		AttributeViewModel.lolcatMode = lolcatMode;
-	}
-
 	@Override
 	public void draw(Graphics2D g) {
 		AttributedString as = this.getAttributedString(g);
@@ -56,8 +48,8 @@ public class AttributeViewModel implements ViewModel {
 	}
 
 	private Color getBackgroundColor(int x, int y, Graphics2D g) {
-		if (!lolcatMode) return BACKGROUND_COLOR;
-		Dimension viewportSize = ((DiagramPanel)EntityRelationMappingEditor.getFrame().getContentPane()).getModel().getRelationBounds().getSize();
+		if (!LolcatAction.getInstance().isLolcatEnabled()) return BACKGROUND_COLOR;
+		Dimension viewportSize = g.getClipBounds().getSize();
 
 		double dx = viewportSize.width;
 		double dy = viewportSize.height;
@@ -67,6 +59,7 @@ public class AttributeViewModel implements ViewModel {
 
 		double lambda = (dx * x + dy * y);
 		double diag_val = Math.sqrt(Math.pow(dx * lambda, 2) + Math.pow(dy * lambda, 2)) / mag;
+		System.out.println(diag_val);
 
 		return Color.getHSBColor((float) diag_val, LOLCAT_SAT, LOLCAT_BRIGHT);
 	}