vilarin commited on
Commit
eee1247
·
verified ·
1 Parent(s): c075c0f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -90,9 +90,9 @@ class ModelWrapper:
90
  else:
91
  raise NotImplementedError()
92
 
93
- noise = noise.to(torch.float16)
94
  print(f'noise: {noise.dtype}')
95
- prompt_embed = prompt_embed.to(torch.float16)
96
  DTYPE = prompt_embed.dtype
97
  print(f'prompt_embed: {DTYPE}')
98
 
@@ -170,7 +170,7 @@ def get_x0_from_noise(sample, model_output, alphas_cumprod, timestep):
170
  return pred_original_sample
171
 
172
  class SDXLTextEncoder(torch.nn.Module):
173
- def __init__(self, model_id, revision, accelerator, dtype=torch.float16):
174
  super().__init__()
175
 
176
  self.text_encoder_one = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder", revision=revision).to(0).to(dtype=dtype)
 
90
  else:
91
  raise NotImplementedError()
92
 
93
+ noise = noise.to(torch.float32)
94
  print(f'noise: {noise.dtype}')
95
+ prompt_embed = prompt_embed.to(torch.float32)
96
  DTYPE = prompt_embed.dtype
97
  print(f'prompt_embed: {DTYPE}')
98
 
 
170
  return pred_original_sample
171
 
172
  class SDXLTextEncoder(torch.nn.Module):
173
+ def __init__(self, model_id, revision, accelerator, dtype=torch.float32):
174
  super().__init__()
175
 
176
  self.text_encoder_one = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder", revision=revision).to(0).to(dtype=dtype)