update code

app.py

@@ -396,13 +396,13 @@ def images_to_video(image_folder, output_video, fps=30):
 
     print(f"✅ High-quality MP4 video has been generated: {output_video}")
 
-@spaces.GPU(duration=100)
+
 def model_define():
     args = get_args()
     set_env(args.seed)
     input_process_model = Process(cfg)
 
-    device = "cuda" 
+    device = "cuda" if torch.cuda.is_available() else "cpu"
     
     weight_dtype = torch.float32
     logging.info(f"Running inference with {weight_dtype}")
@@ -440,18 +440,8 @@ def model_define():
     base_coff = torch.from_numpy(base_coff).float()
     Faceverse = Faceverse_manager(device=device, base_coeff=base_coff)
 
-    controlnet_path = './pretrained_model/control'
-    controlnet = ControlNetModel.from_pretrained(
-        controlnet_path, torch_dtype=torch.float16
-    )
-    sd_path =  './pretrained_model/sd21'
-    pipeline_sd = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
-        sd_path, torch_dtype=torch.float16,
-        use_safetensors=True, controlnet=controlnet, variant="fp16"
-    ).to(device)
-
     return motion_aware_render_model, sample_steps, DiT_model, \
-        vae_triplane, image_encoder, dinov2, dino_img_processor, clip_image_processor, triplane_std, triplane_mean, ws_avg, Faceverse, device, input_process_model,pipeline_sd
+        vae_triplane, image_encoder, dinov2, dino_img_processor, clip_image_processor, triplane_std, triplane_mean, ws_avg, Faceverse, device, input_process_model
 
 
 def duplicate_batch(tensor, batch_size=2):
@@ -460,11 +450,8 @@ def duplicate_batch(tensor, batch_size=2):
     return tensor.repeat(batch_size, *([1] * (tensor.dim() - 1)))  # 复制 batch 维度
 
 
-@torch.inference_mode()
 @spaces.GPU(duration=200)
 def avatar_generation(items, save_path_base, video_path_input, source_type, is_styled, styled_img):
-
-
     """
     Generate avatars from input images.
 
@@ -491,7 +478,15 @@ def avatar_generation(items, save_path_base, video_path_input, source_type, is_s
     exp_img_base_dir = os.path.join(target_path, 'images512x512')
     motion_base_dir = os.path.join(target_path, 'motions')
     label_file_test = os.path.join(target_path, 'images512x512/dataset_realcam.json')
-
+    render_model =render_model.to(device)
+    image_encoder = image_encoder.to(device)
+    vae_triplane = vae_triplane.to(device)
+    dinov2 = dinov2.to(device)
+    Faceverse = Faceverse.to(device)
+    clip_image_processor = clip_image_processor.to(device)
+    dino_img_processor = dino_img_processor.to(device)
+    ws_avg = ws_avg.to(device)
+    DiT_model = DiT_model.to(device)
     if source_type == 'example':
         input_img_fvid = './demo_data/source_img/img_generate_different_domain/coeffs/demo_imgs'
         input_img_motion = './demo_data/source_img/img_generate_different_domain/motions/demo_imgs'
@@ -658,6 +653,7 @@ def style_transfer(processed_image, style_prompt, cfg, strength, save_base):
     🎭 这个函数用于风格转换
     ✅ 你可以在这里填入你的风格化代码
     """
+    pipeline_sd =pipeline_sd.to(device)
     src_img_pil = Image.open(processed_image)
     img_name = os.path.basename(processed_image)
     save_dir = os.path.join(save_base, 'style_img')
@@ -1003,7 +999,15 @@ if __name__ == '__main__':
     image_folder = "./demo_data/source_img/img_generate_different_domain/images512x512/demo_imgs"
     example_img_names = os.listdir(image_folder)
     render_model, sample_steps, DiT_model, \
-        vae_triplane, image_encoder, dinov2, dino_img_processor, clip_image_processor, std, mean, ws_avg, Faceverse, device, input_process_model, pipeline_sd = model_define()
-
+        vae_triplane, image_encoder, dinov2, dino_img_processor, clip_image_processor, std, mean, ws_avg, Faceverse, device, input_process_model = model_define()
+    controlnet_path = './pretrained_model/control'
+    controlnet = ControlNetModel.from_pretrained(
+        controlnet_path, torch_dtype=torch.float16
+    )
+    sd_path =  './pretrained_model/sd21'
+    pipeline_sd = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+        sd_path, torch_dtype=torch.float16,
+        use_safetensors=True, controlnet=controlnet, variant="fp16"
+    ).to(device)
     demo_cam = False
     launch_gradio_app()