Tianshuo-Xu commited on
Commit
46e86e6
·
1 Parent(s): 0108eaf

fix(zerogpu): sync internal device pointer to cuda for noise generation

Browse files
Files changed (1) hide show
  1. app.py +10 -4
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
- gen.model.to("cuda")
284
- gen.clip.to("cuda")
285
- gen.t5.to("cuda")
286
- gen.vae.to("cuda")
 
 
 
 
 
 
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.")