Spaces:
Sleeping
Sleeping
Tianshuo-Xu commited on
Commit ·
46e86e6
1
Parent(s): 0108eaf
fix(zerogpu): sync internal device pointer to cuda for noise generation
Browse files
app.py
CHANGED
|
@@ -280,10 +280,16 @@ def run_generation(text, font, author, num_steps, start_seed, num_images):
|
|
| 280 |
logger.info("Models are already globally initialized and managed by ZeroGPU.")
|
| 281 |
gen = generator
|
| 282 |
# ZeroGPU automatically maps these to the acquired GPU during execution.
|
| 283 |
-
|
| 284 |
-
|
| 285 |
-
gen.
|
| 286 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 287 |
|
| 288 |
# Step 2: Since we reverted to bf16 load to avoid PyTorch native dtype mix issues, skip wrapping
|
| 289 |
logger.info("Model weights decompressed to bfloat16 upon load. Skipping dynamic quantization to ensure stability.")
|
|
|
|
| 280 |
logger.info("Models are already globally initialized and managed by ZeroGPU.")
|
| 281 |
gen = generator
|
| 282 |
# ZeroGPU automatically maps these to the acquired GPU during execution.
|
| 283 |
+
# We must also correctly update internal Python attributes so runtime-generated latents go to GPU.
|
| 284 |
+
target_device = torch.device("cuda")
|
| 285 |
+
gen.device = target_device
|
| 286 |
+
if hasattr(gen, "sampler") and gen.sampler is not None:
|
| 287 |
+
gen.sampler.device = target_device
|
| 288 |
+
|
| 289 |
+
gen.model.to(target_device)
|
| 290 |
+
gen.clip.to(target_device)
|
| 291 |
+
gen.t5.to(target_device)
|
| 292 |
+
gen.vae.to(target_device)
|
| 293 |
|
| 294 |
# Step 2: Since we reverted to bf16 load to avoid PyTorch native dtype mix issues, skip wrapping
|
| 295 |
logger.info("Model weights decompressed to bfloat16 upon load. Skipping dynamic quantization to ensure stability.")
|