vilarin commited on
Commit
20f6444
·
verified ·
1 Parent(s): cb054c0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -123,7 +123,8 @@ class ModelWrapper:
123
 
124
  add_time_ids = self.build_condition_input(height, width).repeat(num_images, 1)
125
 
126
- noise = torch.randn(num_images, 4, height // self.vae_downsample_ratio, width // self.vae_downsample_ratio, generator=generator).to(device="cuda")
 
127
  noise = noise.to(torch.float16)
128
 
129
  prompt_inputs = self._encode_prompt(prompt)
 
123
 
124
  add_time_ids = self.build_condition_input(height, width).repeat(num_images, 1)
125
 
126
+ noise = torch.randn(num_images, 4, height // self.vae_downsample_ratio, width // self.vae_downsample_ratio, generator=generator).to(device="cuda",dtype=torch.float16)
127
+
128
  noise = noise.to(torch.float16)
129
 
130
  prompt_inputs = self._encode_prompt(prompt)