load_model

app.py

@@ -11,10 +11,14 @@ import random
 import spaces
 
 pipe = None
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = None
+torch_dtype = None
 
 def load_model():
-  global device
+  global pipe, device, torch_dtype
+  device = "cuda" if torch.cuda.is_available() else "cpu"
+  torch_dtype = torch.float16 if device == "cuda" else torch.float32
+
   pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
       "yeq6x/animagine_position_map",
       controlnet=ControlNetModel.from_pretrained("yeq6x/Image2PositionColor_v3"),
@@ -23,8 +27,6 @@ def load_model():
   
   return pipe
 
-pipe = load_model()
-
 def convert_pil_to_opencv(pil_image):
   return np.array(pil_image)
 
@@ -139,9 +141,8 @@ def outpaint_image(image):
 
 @spaces.GPU
 def predict_image(cond_image, prompt, negative_prompt):
-  print("Processing...")
+  print("predict position map")
   global pipe
-  print("Processing...")
   generator = torch.Generator()
   generator.manual_seed(random.randint(0, 2147483647))
   image = pipe(
@@ -160,6 +161,8 @@ def predict_image(cond_image, prompt, negative_prompt):
   
   return image
 
+load_model()
+
 # Gradioアプリケーション
 with gr.Blocks() as demo:
   gr.Markdown("## Position Map Visualizer")