Fix handling of figures and update demo
Finally the right fix for plotting.