Sudden Tony Change

In Made in America (one of the best episodes of the absolute best thing ever made of course) Tony kind of goes from this super dark descent into madness guy but then suddenly is a nice guy dad playing journey and eating onion rings, I feel the transition is off somehow, like it’s too extreme and just made so that last scene can exist. Also we never see or hear about that diner in the whole show besides the end of that episode but they are like oh year we eat here all the time. Maybe that’s the point somehow? It’s all a fantasy?