bfloat

app.py

@@ -160,9 +160,8 @@ def encode_image(pil_image):
 
 
 ## --- Model Loading --- ##
-# device = "cuda" if torch.cuda.is_available() else "cpu"
-device = "cpu"
-dtype = torch.float16 if device == "cuda" else torch.float32
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.bfloat16
 print(f"Using device: {device}")
 
 # controlnet
@@ -281,7 +280,7 @@ def infer(
     # Blur mask
     # data["mask_img"] = data["mask_img"].filter(ImageFilter.GaussianBlur(15))
 
-    # denoise with tqdm
+    # denoise
     outpaint_result = pipe(
         prompt=generated_text,
         negative_prompt="NSFW, (word:1.5), watermark, blurry, missing body, amputation, mutilation",
@@ -292,21 +291,8 @@ def infer(
         guidance_scale=cfg,
         generator=generator,
         repeat_time=repeat_time,
-        progress=progress,
     ).images[0]
 
-    # outpaint_result = pipe(
-    #     prompt=generated_text,
-    #     negative_prompt="NSFW, (word:1.5), watermark, blurry, missing body, amputation, mutilation",
-    #     image=data["input_img"],
-    #     mask_image=data["mask_img"],
-    #     control_image=data["control_img"],
-    #     num_inference_steps=num_inference_steps,
-    #     guidance_scale=cfg,
-    #     generator=generator,
-    #     repeat_time=repeat_time,
-    # ).images[0]
-
     # save result
     outpaint_result = np.array(outpaint_result)
     input_img = np.array(data["input_img"])