update code

app.py

@@ -305,7 +305,10 @@ def prepare_working_dir(dir, style):
 def launch_pretrained():
 
     from huggingface_hub import snapshot_download
-
+    os.system("pip uninstall torch")
+    os.system("pip uninstall torchvision")
+    os.system("pip install https://download.pytorch.org/whl/cu121/torch-2.4.1%2Bcu121-cp310-cp310-linux_x86_64.whl")
+    os.system("pip install https://download.pytorch.org/whl/cu121/torchvision-0.19.1%2Bcu121-cp310-cp310-linux_x86_64.whl")
     snapshot_download(
         repo_id="KumaPower/AvatarArtist",
         repo_type="model",
@@ -954,11 +957,15 @@ if __name__ == '__main__':
     if _get_output(["nvcc", "--version"]) is None:
         logging.info("Installing CUDA toolkit...")
         install_cuda_toolkit()
+        logging.info("install CUDA: %s" % _get_output(["nvcc", "--version"]))
+
     else:
         logging.info("Detected CUDA: %s" % _get_output(["nvcc", "--version"]))
 
 
     launch_pretrained()
+
+
     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, \