test

app.py

@@ -76,8 +76,9 @@ def enhance_face(img, face_helper, has_aligned, only_center_face=False, paste_ba
         cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
 
         dummy_x = torch.zeros_like(cropped_face_t)
-        output = generate_reconstructions(pmrf, dummy_x, cropped_face_t, None, 25, device)
-        restored_face = tensor2img(output.squeeze(0), rgb2bgr=True, min_max=(0, 1))
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            output = generate_reconstructions(pmrf, dummy_x, cropped_face_t, None, 25, device)
+        restored_face = tensor2img(output.to(torch.float32).squeeze(0), rgb2bgr=True, min_max=(0, 1))
         # restored_face = cropped_face
 
         restored_face = restored_face.astype('uint8')

arch/hourglass/image_transformer_v2.py

@@ -92,7 +92,7 @@ def linear_geglu(x, weight, bias=None):
     if bias is not None:
         x = x + bias
     x, gate = x.chunk(2, dim=-1)
-    return x.clone() * F.gelu(gate.clone())
+    return x * F.gelu(gate)
 
 
 def rms_norm(x, scale, eps):