Pygame key-presses don't work on 2nd run of cell in Jupyter Lab

Question

When I run the code below after a kernel restart in Jupyter Lab if I hit the space bar, a highlighted 'aaa' will result (as it should).

But, if I ever have to re-run the same cell again (without a kernel restart), a space bar hit will just result in Jupyter Lab scrolling downwards and it being ignored by pygame although the program in every other respect is still functioning

Any ideas on how I can get the key-presses working every time I run that cell in Jupyter Lab?

thanks

offset = 2130

#!/usr/bin/env python3

import argparse
import os
import queue
import sounddevice as sd
import vosk
import sys
import ast
import time
import colored
import pygame


start = time.time()
pygame.init()

q = queue.Queue()



def int_or_str(text):
    """Helper function for argument parsing."""
    try:
        return int(text)
    except ValueError:
        return text

def callback(indata, frames, time, status):
    """This is called (from a separate thread) for each audio block."""
    if status:
        print(status, file=sys.stderr)
    q.put(bytes(indata))

parser = argparse.ArgumentParser(add_help=False)
parser.add_argument(
    '-l', '--list-devices', action='store_true',
    help='show list of audio devices and exit')
args, remaining = parser.parse_known_args()
if args.list_devices:
    print(sd.query_devices())
    parser.exit(0)
parser = argparse.ArgumentParser(
    description=__doc__,
    formatter_class=argparse.RawDescriptionHelpFormatter,
    parents=[parser])
parser.add_argument(
    '-f', '--filename', type=str, metavar='FILENAME',
    help='audio file to store recording to')
parser.add_argument(
    '-m', '--model', type=str, metavar='MODEL_PATH',
    help='Path to the model')
parser.add_argument(
    '-d', '--device', type=int_or_str,
    help='input device (numeric ID or substring)')
parser.add_argument(
    '-r', '--samplerate', type=int, help='sampling rate')
args = parser.parse_args(remaining)

try:
    if args.model is None:
        args.model = "model"
    if not os.path.exists(args.model):
        print ("Please download a model for your language from https://alphacephei.com/vosk/models")
        print ("and unpack as 'model' in the current folder.")
        parser.exit(0)
    if args.samplerate is None:
        device_info = sd.query_devices(args.device, 'input')
        # soundfile expects an int, sounddevice provides a float:
        args.samplerate = int(device_info['default_samplerate'])

    model = vosk.Model(args.model)

    if args.filename:
        dump_fn = open(args.filename, "wb")
    else:
        dump_fn = None

    with sd.RawInputStream(samplerate=args.samplerate, blocksize = 8000, device=args.device, dtype='int16',
                            channels=1, callback=callback):
            print('#' * 80)
            print('Press Ctrl+C to stop the recording')
            print('#' * 80)
            
            print(f'offset = {offset}')

            rec = vosk.KaldiRecognizer(model, args.samplerate)
            xxx = True
            while xxx==True:
                
                data = q.get()
                if rec.AcceptWaveform(data):
            
                    yag = rec.Result()
                    yag3 = ast.literal_eval(yag)
                    
                    sec = time.time()-start + offset
                    ty_res = time.gmtime(sec)
                    res = time.strftime("%H:%M:%S",ty_res)
    
                    try:
                        print(f"{res} - {yag3['text']}")
                    except:
                        print(f"{res} - oops")
                else:
                    pass

                if dump_fn is not None:
                    dump_fn.write(data)
                for event in pygame.event.get():
                    if event.type == pygame.KEYDOWN:
                        if event.key == pygame.K_RETURN:
                            print('\nDone') 
                            xxx = False
#                             parser.exit(0)
                            pygame.quit()
                            
 
                            
                        if event.key == pygame.K_SPACE:
                            print(f"{colored.bg(226)} aaa {colored.bg(231)} ")   

except KeyboardInterrupt:
    print('\nDone')  
    pygame.quit()
    parser.exit(0)

except Exception as e:
    parser.exit(type(e).__name__ + ': ' + str(e))
    pygame.quit()