fix: device sync issue

app.py

@@ -125,9 +125,7 @@ def get_model(name):
  subfolder="unet",
  torch_dtype=torch.float16,
  )
- if torch.cuda.is_available():
- unet.to("cuda")
-
+ unet.to("cuda")
  unet_cache[name] = unet
  lora_cache[name] = LoRANetwork(lora_cache[base_name].text_encoder_loras, unet)
 
@@ -135,6 +133,9 @@ def get_model(name):
  g_lora = lora_cache[name]
  g_unet.set_attn_processor(CrossAttnProcessor())
  g_lora.reset()
+ if torch.cuda.is_available():
+ g_unet.to("cuda")
+ g_lora.to("cuda")
  return g_unet, g_lora, models[keys.index(name)][2]
 
 # precache on huggingface