Sachverhalt: A. Mit Urteil vom 15. Juni 2016 verurteilte das Bezirksgericht Zürich A._ wegen mehrfacher, teilweise versuchter Gewalt und Drohung gegen Behörden und Beamte (Art. 285 Ziff. 1 StGB, teilweise i.V.m. Art. 22 Abs. 1 StGB) zu einer Freiheitsstrafe von acht Monaten, wovon bis und mit Urteilsdatum 231 Tage durch Haft und vorzeitigen Strafantritt erstanden waren. Weiter wurde eine stationäre therapeutische Massnahme (Art. 59 StGB; Behandlung von psychischen Störungen) angeordnet und die Freiheitsstrafe zu diesem Zweck aufgeschoben. Gegen dieses Urteil erhob A._ am 24. Juni 2016 Berufung. B. Mit Beschluss vom 6. Juli 2016 wies das Bezirksgericht das Haftentlassungsgesuch von A._ ab. Es entliess ihn aus dem vorzeitigen Strafvollzug und versetze ihn sogleich in Sicherheitshaft. Diese wurde bis zum 6. Oktober 2016 befristet. Die dagegen von A._ am 20. Juli 2016 erhoben Beschwerde wies das Obergericht des Kantons Zürich, III. Strafkammer, mit Beschluss vom 29. Juli 2016 ab. Seit anfangs August 2016 befindet sich A._ im vorzeitigen stationären Massnahmenvollzug in der Klinik N._. C. A._ erhebt Beschwerde in Strafsachen. Er beantragt die Aufhebung des angefochtenen Beschlusses und die umgehende Entlassung aus dem vorzeitigen Massnahmenvollzug. Als Ersatzmassnahme sei eine ambulante psychiatrische Behandlung anzuordnen. Das Obergericht und das Bezirksgericht haben auf eine Vernehmlassung verzichtet. Die Staatsanwaltschaft beantragt die Abweisung der Beschwerde. Der Beschwerdeführer hat auf eine Stellungnahme verzichtet.