Mixing signal with noise in frequency space (STFT)

Question

I am working on a preprocessing function which is supposed to mix two audio signals clean and noise in the context of neural speech recognition.

For this I would like to augment the clean sample and mix it later with the noise sample.

Now, for efficiency reasons I would like to mix those to samples in frequency-space, as opposed to mixing them in time-space.

I wrote some Python code a while ago which mixes to signals randomly based on a given distribution and makes sure, that the noise sample does not get "too loud" relative to the clean signal:

# Calculate power of signals
s_pow = np.sum(np.abs(signal_waveform) ** 2.0) / len(signal_waveform)
n_pow = np.sum(np.abs(noise_waveform) ** 2.0) / len(noise_waveform)

# Draw SNR-value from given distribution
snr = eval('self.rnd.%s' % self.distribution)

# Calculate the scaling factor from SNR = 10 * log(P_s / P_n)
k = (s_pow / n_pow) * 10 ** (-snr / 10)
new_noise_waveform = np.sqrt(k) * noise_waveform

# Apply noise
final_signal_waveform = signal_waveform + new_noise_waveform

s = np.iinfo(signal_waveform.dtype).max / np.max(final_signal_waveform)
# Scale signal down as a whole if necessary to avoid clipping
if s < 1.0: 
  final_signal_waveform = final_signal_waveform * s

However, I was wondering how this could be achieved in frequency space. Here, I would like to augment the clean sample first in frequency-space and then mix it with the stft of the noise sample:

clean_stft = spectrum.stft(clean)
noise_stft = spectrum.stft(noise)

augmented_stft = spectrum.augment(clean_stft)  # e.g. time_stretch()

# How can this be done?
mixed = spectrum.mix(augmented_stft, noise_stft)

mel = spectrum.mel_features(mixed)

The reason is efficiency here. I don't want to convert everything back into time-space only having to run stft calculations again moments later.

So: How can I mix two stft spectrograms s.t. I can also control the relative loudness?