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