diff --git a/src/trezor/layout.py b/src/trezor/layout.py
index 2190ee63e1..0f4dcd85a4 100644
--- a/src/trezor/layout.py
+++ b/src/trezor/layout.py
@@ -1,7 +1,9 @@
+import sys
+
 _new_layout = None
 _current_layout = None
 
-def change_layout(layout):
+def change(layout):
     global _new_layout
 
     print("Changing layout to %s" % layout)
@@ -9,27 +11,26 @@ def change_layout(layout):
 
     yield _current_layout.throw(StopIteration())
 
-def set_main_layout(main_layout):
+def set_main(main_layout):
     global _new_layout
     global _current_layout
 
-    layout = main_layout
+    _current_layout = main_layout
     while True:
         try:
-            _current_layout = layout()
-            layout = yield from _current_layout
+            _current_layout = yield from _current_layout
         except Exception as e:
-            print("Layout thrown exception %s" % str(e))
+            sys.print_exception(e)
             _current_layout = main_layout
             continue
 
         if _new_layout != None:
             print("Switching to new layout %s" % _new_layout)
-            layout = _new_layout
+            _current_layout = _new_layout
             _new_layout = None
 
-        elif layout == None:
+        elif not callable(_current_layout):
             print("Switching to main layout %s" % main_layout)
-            layout = main_layout
+            _current_layout = main_layout
         else:
-            print("Switching to proposed layout %s" % layout)
+            print("Switching to proposed layout %s" % _current_layout)