update code

app.py

@@ -56,19 +56,7 @@ else:
 
 import torch
 
-# 设置当前默认 GPU 设备（推荐在 CUDA 初始化前设置）
-torch.cuda.set_device(0)
-
-# 显式初始化 CUDA（通常是可选的，但在多线程中有助于避免问题）
-torch.cuda.init()
-
-# 测试
-print("CUDA available:", torch.cuda.is_available())
-print("Current device:", torch.cuda.current_device())
-print("Device name:", torch.cuda.get_device_name(0))
-print("CUDA_HOME =", os.environ.get("CUDA_HOME"))
-from torch.utils.cpp_extension import CUDA_HOME
-print("CUDA_HOME from PyTorch:", CUDA_HOME)
+
 import argparse
 import json
 import random
@@ -264,6 +252,7 @@ def generate_samples(DiT_model, cfg_scale, sample_steps, clip_feature, dino_feat
 
 
 def load_motion_aware_render_model(ckpt_path, device):
+    
     """Load the motion-aware render model from a checkpoint."""
     logging.info("Loading motion-aware render model...")
     with dnnlib.util.open_url(ckpt_path, 'rb') as f:
@@ -407,13 +396,14 @@ 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" if torch.cuda.is_available() else "cpu"
+    device = "cuda" 
+    
     weight_dtype = torch.float32
     logging.info(f"Running inference with {weight_dtype}")
 
@@ -450,8 +440,18 @@ 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
+        vae_triplane, image_encoder, dinov2, dino_img_processor, clip_image_processor, triplane_std, triplane_mean, ws_avg, Faceverse, device, input_process_model,pipeline_sd
 
 
 def duplicate_batch(tensor, batch_size=2):
@@ -463,6 +463,8 @@ def duplicate_batch(tensor, batch_size=2):
 @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.
 
@@ -650,7 +652,7 @@ def process_image(input_image, source_type, is_style, save_dir):
         imge_dir = os.path.join(save_dir, 'processed_img/dataset/images512x512/input_image', img_name)
         return imge_dir, source_type  # 这里替换成 处理用户上传图片的逻辑
 
-
+@spaces.GPU(duration=100)
 def style_transfer(processed_image, style_prompt, cfg, strength, save_base):
     """
     🎭 这个函数用于风格转换
@@ -1001,15 +1003,7 @@ 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 = 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)
+        vae_triplane, image_encoder, dinov2, dino_img_processor, clip_image_processor, std, mean, ws_avg, Faceverse, device, input_process_model, pipeline_sd = model_define()
+
     demo_cam = False
     launch_gradio_app()