Update main.py

fix bugs !!!