Update app.py

app.py

@@ -91,9 +91,6 @@ class ModelWrapper:
         
         DTYPE = prompt_embed.dtype
         print(DTYPE)
-        print(type(noise))
-        print(type(current_timesteps))
-        print(type(unet_added_conditions))
         
         for constant in all_timesteps:
             current_timesteps = torch.ones(len(prompt_embed), device="cuda", dtype=torch.long) * constant
@@ -124,7 +121,7 @@ class ModelWrapper:
 
         add_time_ids = self.build_condition_input(height, width).repeat(num_images, 1)
 
-        noise = torch.randn(num_images, 4, height // self.vae_downsample_ratio, width // self.vae_downsample_ratio, generator=generator).to(device="cuda", dtype=self.DTYPE) 
+        noise = torch.randn(num_images, 4, height // self.vae_downsample_ratio, width // self.vae_downsample_ratio, generator=generator).to(device="cuda", dtype=float16) 
 
         prompt_inputs = self._encode_prompt(prompt)
         
@@ -161,7 +158,7 @@ def get_x0_from_noise(sample, model_output, alphas_cumprod, timestep):
     return pred_original_sample
 
 class SDXLTextEncoder(torch.nn.Module):
-    def __init__(self, model_id, revision, accelerator, dtype=torch.float32):
+    def __init__(self, model_id, revision, accelerator, dtype=torch.float16):
         super().__init__()
 
         self.text_encoder_one = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder", revision=revision).to(0).to(dtype=dtype)