InvalidArgumentError when using GPU Colab + Mixed Precision with input_image

[costiash]

Hello, first and foremost, thank you for this fantastic repository and the provided colab demonstrations! When I attempted the GPU Colab + Mixed Precision and tried to send an input image argument to the generator, I received the following error: ` InvalidArgumentError Traceback (most recent call last)

in 7 temperature=1, 8 batch_size=1, ----> 9 input_image="/content/gen.png" 10 ) 4 frames /usr/local/lib/python3.7/dist-packages/stable_diffusion_tf/stable_diffusion.py in generate(self, prompt, batch_size, num_steps, unconditional_guidance_scale, temperature, seed, input_image, input_image_strength) 74 input_img_noise_t = timesteps[ int(len(timesteps)*input_image_strength) ] 75 latent, alphas, alphas_prev = self.get_starting_parameters( ---> 76 timesteps, batch_size, seed , input_image=input_image, input_img_noise_t=input_img_noise_t 77 ) 78 /usr/local/lib/python3.7/dist-packages/stable_diffusion_tf/stable_diffusion.py in get_starting_parameters(self, timesteps, batch_size, seed, input_image, input_img_noise_t) 159 else: 160 latent = self.encoder(input_image[None]) --> 161 latent = self.add_noise(latent, input_img_noise_t) 162 latent = tf.repeat(latent , batch_size , axis=0) 163 return latent, alphas, alphas_prev /usr/local/lib/python3.7/dist-packages/stable_diffusion_tf/stable_diffusion.py in add_noise(self, x, t) 108 sqrt_one_minus_alpha_prod = (1 - _ALPHAS_CUMPROD[t]) ** 0.5 109 --> 110 return sqrt_alpha_prod * x + sqrt_one_minus_alpha_prod * noise 111 112 def timestep_embedding(self, timesteps, dim=320, max_period=10000): /usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py in error_handler(*args, **kwargs) 151 except Exception as e: 152 filtered_tb = _process_traceback_frames(e.__traceback__) --> 153 raise e.with_traceback(filtered_tb) from None 154 finally: 155 del filtered_tb /usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py in raise_from_not_ok_status(e, name) 7207 def raise_from_not_ok_status(e, name): 7208 e.message += (" name: " + name if name is not None else "") -> 7209 raise core._status_to_exception(e) from None # pylint: disable=protected-access 7210 7211 InvalidArgumentError: cannot compute AddV2 as input #1(zero-based) was expected to be a half tensor but is a float tensor [Op:AddV2] ` first. I ran the StableDiffusion generator instantiation and created the first picture, which worked as expected. Then I tried running the generator again, passing the picture made in the previous run as the 'input image' option, but I got the above issue. I also adjusted the installation in the Colab demo to be from the most recent commit of this repo: `pip install git+https://github.com/divamgupta/stable-diffusion-tensorflow --upgrade --quiet` Thanks a lot in advance!

[costiash]

Looks like I have managed to find a work around. feels somewhat like weird solution but it worked. I have added to add_noise method the following:

    def add_noise(self, x , t ):
        batch_size,w,h = x.shape[0] , x.shape[1] , x.shape[2]
        noise = tf.random.normal((batch_size,w,h,4))
        sqrt_alpha_prod = _ALPHAS_CUMPROD[t] ** 0.5
        sqrt_one_minus_alpha_prod = (1 - _ALPHAS_CUMPROD[t]) ** 0.5

        sqrt_alpha_prod_mul_x = sqrt_alpha_prod * x
        sqrt_alpha_prod_mul_x = tf.cast(sqrt_alpha_prod_mul_x, dtype=tf.float32)

        sqrt_one_minus_alpha_prod_mul_noise = sqrt_one_minus_alpha_prod * noise
        sqrt_one_minus_alpha_prod_mul_noise = tf.cast(sqrt_one_minus_alpha_prod_mul_noise, dtype=tf.float32)

        return  sqrt_alpha_prod_mul_x + sqrt_one_minus_alpha_prod_mul_noise