Improve edit screen