If you are using diatonic chords - that, by definition, use the notes of the home scale - it would be perverse to use anything other than those notes as the basis of your playing. (Though note that some perversity can be interesting!)
I have never quite understood why some people insist on renaming the scale according to WHICH diatonic chord is being played - D Dorian for a Dm chord, E Phrygian for an Em chord - when they all contain the same notes. What's wrong with admitting it's all the scale of C major? But they want to learn 7 different scales instead of just one.
When a chromatic chord is used, you can think of it as being diatonic in some OTHER key, or as a splash of colour in the original key. Your choice may depend on how quickly returns to the home tonality.
Don't rely too much on 'theory'. Improvisation should grow from playing lots of music (you DO read music, don't you?) from transcribing and copying other players' soloing, not from learning a set of rules. Don't improvise until you can PLAY!