Update app.py

app.py

@@ -20,22 +20,25 @@ config_file = hf_hub_download(
 config = ControlNetModel_Union.load_config(config_file)
 controlnet_model = ControlNetModel_Union.from_config(config)
 
-
 # Load the state dictionary
 model_file = hf_hub_download(
     "xinsir/controlnet-union-sdxl-1.0",
     filename="diffusion_pytorch_model_promax.safetensors",
 )
-state_dict = load_state_dict(model_file)  # Define state_dict
+state_dict = load_state_dict(model_file)
 
-# Then extract the keys from the state_dict
+# Extract the keys from the state_dict
 loaded_keys = list(state_dict.keys())
 
-# Then call the method with the additional parameter
-model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
+# Call the method and store all returns in a variable
+result = ControlNetModel_Union._load_pretrained_model(
     controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0", loaded_keys
 )
 
+# Use the first element from the result
+model = result[0]
+model = model.to(device="cuda", dtype=torch.float16)
+
 
 vae = AutoencoderKL.from_pretrained(
     "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16