dynamically get VAE mapping before pairing with ControlNet

app.py

@@ -79,11 +79,13 @@ pipe = QwenImageControlNetInpaintPipeline.from_pretrained(
     device_map="balanced"
 )
 
-print("Initial device map:", pipe.hf_device_map)
+pipe_device_map = pipe.hf_device_map
+print("Initial device map:", pipe_device_map)
 # Expected output: {'transformer': 0, 'text_encoder': 1, 'vae': 2}
 
 # Move the controlnet to the same device as the VAE (cuda:2)
-vae_device = "cuda:2"  # This is where the VAE is in the balanced config
+vae_device = pipe_device_map['vae']
+vae_device = f"cuda:{vae_device}"  # This is where the VAE is in the balanced config
 controlnet = QwenImageControlNetModel.from_pretrained(
     controlnet_model, 
     torch_dtype=torch.bfloat16