gpu

batch_generator.py

@@ -23,7 +23,7 @@ import numpy as np
 import gc
 
 # Global device configuration
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 dtype = torch.float16
 
 # Initialize global generator
@@ -33,7 +33,7 @@ generator = torch.Generator()
 print("Loading VAE...")
 vae = AutoencoderKL.from_pretrained(
     "madebyollin/sdxl-vae-fp16-fix", torch_dtype=dtype, use_safetensors=True
-).to(device)
+).to("cuda")
 
 print("Loading base pipeline...")
 pipe = StableDiffusionXLPipeline.from_pretrained(
@@ -42,7 +42,7 @@ pipe = StableDiffusionXLPipeline.from_pretrained(
     torch_dtype=dtype,
     variant="fp16",
     use_safetensors=True,
-)
+).to("cuda")
 
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
 
@@ -58,7 +58,7 @@ openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
 print("Loading T2I adapter...")
 adapter = T2IAdapter.from_pretrained(
     "TencentARC/t2i-adapter-openpose-sdxl-1.0", torch_dtype=dtype
-)
+).to("cuda")
 
 print("Loading adapter pipeline...")
 posepipe = StableDiffusionXLAdapterPipeline.from_pretrained(
@@ -68,8 +68,7 @@ posepipe = StableDiffusionXLAdapterPipeline.from_pretrained(
     torch_dtype=dtype,
     variant="fp16",
     use_safetensors=True,
-    attn_implementation="xformers",
-)
+).to("cuda")
 
 
 posepipe.scheduler = UniPCMultistepScheduler.from_config(posepipe.scheduler.config)