update

app.py

@@ -26,20 +26,23 @@ with gr.Blocks(title="Anime AI Detect Fucker Demo", theme="dark") as demo:
     gr.HTML('<a href="https://github.com/7eu7d7/anime-ai-detect-fucker">github repo</a>')
 
     with gr.Row():
-        eps = gr.Slider(label="eps (Noise intensity)", minimum=1, maximum=16, step=1, value=1)
-        step_size = gr.Slider(label="Noise step size", minimum=0.001, maximum=16, step=0.001, value=0.136)
-    with gr.Row():
-        steps = gr.Slider(label="step count", minimum=1, maximum=100, step=1, value=20)
-        model_name = gr.Dropdown(label="attack target",
-                                 choices=["auto", "human", "ai"],
-                                 value="auto", show_label=True)
-
-    input_image = gr.Image(label="Clean Image", type="pil")
-
-    atk_btn = gr.Button("Attack")
-
-    with gr.Column():
-        output_image = gr.Image(label="Attacked Image")
+        with gr.Column():
+            with gr.Row():
+                eps = gr.Slider(label="eps (Noise intensity)", minimum=1, maximum=16, step=1, value=1)
+                step_size = gr.Slider(label="Noise step size", minimum=0.001, maximum=16, step=0.001, value=0.136)
+            with gr.Row():
+                steps = gr.Slider(label="step count", minimum=1, maximum=100, step=1, value=20)
+                model_name = gr.Dropdown(label="attack target",
+                                         choices=["auto", "human", "ai"],
+                                         interactive=True,
+                                         value="auto", show_label=True)
+
+            input_image = gr.Image(label="Clean Image", type="pil")
+
+            atk_btn = gr.Button("Attack")
+
+        with gr.Column():
+            output_image = gr.Image(label="Attacked Image")
 
     atk_btn.click(fn=do_attack,
                       inputs=[input_image, eps, step_size, steps],

attack.py

@@ -35,7 +35,7 @@ class Attacker:
 
         print('正在加载模型...')
         self.feature_extractor = BeitFeatureExtractor.from_pretrained('saltacc/anime-ai-detect')
-        self.model = BeitForImageClassification.from_pretrained('saltacc/anime-ai-detect').cuda()
+        self.model = BeitForImageClassification.from_pretrained('saltacc/anime-ai-detect').to(device)
         print('加载完毕')
 
         if args.target=='ai': #攻击成被识别为AI
@@ -43,8 +43,8 @@ class Attacker:
         elif args.target=='human':
             self.target = torch.tensor([0]).to(device)
 
-        dataset_mean_t = torch.tensor([0.5, 0.5, 0.5]).view(1, -1, 1, 1).cuda()
-        dataset_std_t = torch.tensor([0.5, 0.5, 0.5]).view(1, -1, 1, 1).cuda()
+        dataset_mean_t = torch.tensor([0.5, 0.5, 0.5]).view(1, -1, 1, 1).to(device)
+        dataset_std_t = torch.tensor([0.5, 0.5, 0.5]).view(1, -1, 1, 1).to(device)
         self.pgd = PGD(self.model, img_transform=(lambda x: (x - dataset_mean_t) / dataset_std_t, lambda x: x * dataset_std_t + dataset_mean_t))
         self.pgd.set_para(eps=(args.eps * 2) / 255, alpha=lambda: (args.step_size * 2) / 255, iters=args.steps)
         self.pgd.set_loss(CrossEntropyLoss())
@@ -58,7 +58,7 @@ class Attacker:
         save_image(img_save, os.path.join(self.args.out_dir, f'{img_name[:img_name.rfind(".")]}_atk.png'))
 
     def attack_(self, image):
-        inputs = self.feature_extractor(images=image, return_tensors="pt")['pixel_values'].cuda()
+        inputs = self.feature_extractor(images=image, return_tensors="pt")['pixel_values'].to(device)
 
         if self.args.target == 'auto':
             with torch.no_grad():