Update app_ctrlx.py

app_ctrlx.py

@@ -9,6 +9,8 @@ from ctrl_x.pipelines.pipeline_sdxl import CtrlXStableDiffusionXLPipeline
 from ctrl_x.utils import *
 from ctrl_x.utils.sdxl import *
 
+import spaces
+
 
 parser = ArgumentParser()
 parser.add_argument("-m", "--model", type=str, default=None)  # Optionally, load model checkpoint from single file
@@ -21,19 +23,19 @@ torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 model_id_or_path = "stabilityai/stable-diffusion-xl-base-1.0"
 refiner_id_or_path = "stabilityai/stable-diffusion-xl-refiner-1.0"
 device = "cuda" if torch.cuda.is_available() else "cpu"
-variant = "fp16" if device == "cuda" else "fp32"
+#variant = "fp16" if device == "cuda" else "fp32"
 
 scheduler = DDIMScheduler.from_config(model_id_or_path, subfolder="scheduler")  # TODO: Support other schedulers
 if args.model is None:
     pipe = CtrlXStableDiffusionXLPipeline.from_pretrained(
-        model_id_or_path, scheduler=scheduler, torch_dtype=torch_dtype, variant=variant, use_safetensors=True
+        model_id_or_path, scheduler=scheduler, torch_dtype=torch_dtype, use_safetensors=True
     )
 else:
     print(f"Using weights {args.model} for SDXL base model.")
     pipe = CtrlXStableDiffusionXLPipeline.from_single_file(args.model, scheduler=scheduler, torch_dtype=torch_dtype)
 refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
     refiner_id_or_path, scheduler=scheduler, text_encoder_2=pipe.text_encoder_2, vae=pipe.vae,
-    torch_dtype=torch_dtype, variant=variant, use_safetensors=True,
+    torch_dtype=torch_dtype, use_safetensors=True,
 )
 
 if torch.cuda.is_available():
@@ -155,7 +157,7 @@ title = """
 </div>
 """
 
-
+@spaces.GPU
 def inference(
     structure_image, appearance_image,
     prompt, structure_prompt, appearance_prompt,