Spaces:
Running
on
Zero
Running
on
Zero
tight-inversion
commited on
Commit
·
9499d13
1
Parent(s):
a11990f
Test
Browse files
app.py
CHANGED
|
@@ -135,6 +135,7 @@ class FluxGenerator:
|
|
| 135 |
x = rearrange(x, "h w c -> 1 c h w") # shape: (1, C, H, W)
|
| 136 |
x = x.to(self.device)
|
| 137 |
# Encode with autocast
|
|
|
|
| 138 |
with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16):
|
| 139 |
x = self.ae.encode(x)
|
| 140 |
|
|
|
|
| 135 |
x = rearrange(x, "h w c -> 1 c h w") # shape: (1, C, H, W)
|
| 136 |
x = x.to(self.device)
|
| 137 |
# Encode with autocast
|
| 138 |
+
self.ae.encoder.to(self.device)
|
| 139 |
with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16):
|
| 140 |
x = self.ae.encode(x)
|
| 141 |
|