<s> import numpy as np <EOL> import viz <EOL> def diffusion_step ( Xmid , t , get_mu_sigma , denoise_sigma , mask , XT , rng ) : <EOL> """<STR_LIT>""" <EOL> mu , sigma = get_mu_sigma ( Xmid , np . array ( [ [ t ] ] ) ) <EOL> if denoise_sigma is not None : <EOL> sigma_new = ( sigma ** - <NUM_LIT:2> + denoise_sigma ** - <NUM_LIT:2> ) ** - <NUM_LIT:0.5> <EOL> mu_new = mu * sigma_new ** <NUM_LIT:2> * sigma ** - <NUM_LIT:2> + XT * sigma_new ** <NUM_LIT:2> * denoise_sigma ** - <NUM_LIT:2> <EOL> sigma = sigma_new <EOL> mu = mu_new <EOL> if mask is not None : <EOL> mu . flat [ mask ] = XT . flat [ mask ] <EOL> sigma . flat [ mask ] = <NUM_LIT:0.> <EOL> Xmid = mu + sigma * rng . normal ( size = Xmid . shape ) <EOL> return Xmid <EOL> def generate_inpaint_mask ( n_samples , n_colors , spatial_width ) : <EOL> """<STR_LIT>""" <EOL> mask = np . zeros ( ( n_samples , n_colors , spatial_width , spatial_width ) , dtype = bool ) <EOL> mask [ : , : , : , spatial_width / <NUM_LIT:2> : ] = True <EOL> return mask . ravel ( ) <EOL> def generate_samples ( model , get_mu_sigma , <EOL> n_samples = <NUM_LIT> , inpaint = False , denoise_sigma = None , X_true = None , <EOL> base_fname_part1 = "<STR_LIT>" , base_fname_part2 = '<STR_LIT>' , <EOL> num_intermediate_plots = <NUM_LIT:4> , seed = <NUM_LIT> ) : <EOL> """<STR_LIT>""" <EOL> rng = np . random . RandomState ( seed ) <EOL> spatial_width = model . spatial_width <EOL> n_colors = model . n_colors <EOL> XT = rng . normal ( size = ( n_samples , n_colors , spatial_width , spatial_width ) ) <EOL> if denoise_sigma is not None : <EOL> XT = X_true + XT * denoise_sigma <EOL> base_fname_part1 += '<STR_LIT>' % denoise_sigma <EOL> if inpaint : <EOL> mask = generate_inpaint_mask ( n_samples , n_colors , spatial_width ) <EOL> XT . flat [ mask ] = X_true . flat [ mask ] <EOL> base_fname_part1 += '<STR_LIT>' <EOL> else : <EOL> mask = None <EOL> if X_true is not None : <EOL> viz . plot_images ( X_true , base_fname_part1 + '<STR_LIT>' + base_fname_part2 ) <EOL> viz . plot_images ( XT , base_fname_part1 + '<STR_LIT>' % model . trajectory_length + base_fname_part2 ) <EOL> Xmid = XT . copy ( ) <EOL> for t in xrange ( model . trajectory_length - <NUM_LIT:1> , <NUM_LIT:0> , - <NUM_LIT:1> ) : <EOL> Xmid = diffusion_step ( Xmid , t , get_mu_sigma , denoise_sigma , mask , XT , rng ) <EOL> if np . mod ( model . trajectory_length - t , <EOL> int ( np . ceil ( model . trajectory_length / ( num_intermediate_plots + <NUM_LIT> ) ) ) ) == <NUM_LIT:0> : <EOL> viz . plot_images ( Xmid , base_fname_part1 + '<STR_LIT>' % t + base_fname_part2 ) <EOL> X0 = Xmid <EOL> viz . plot_images ( X0 , base_fname_part1 + '<STR_LIT>' % <NUM_LIT:0> + base_fname_part2 ) </s>
