update code

app.py

@@ -1,7 +1,59 @@
 import os
+import subprocess
 import sys
 import warnings
 import logging
+# Configure logging settings
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s"
+)
+def _get_output(cmd):
+    try:
+        return subprocess.check_output(cmd).decode("utf-8")
+    except Exception as ex:
+        logging.exception(ex)
+
+    return None
+def install_cuda_toolkit():
+    CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run"
+    CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
+
+    print(f"[INFO] Downloading CUDA Toolkit from {CUDA_TOOLKIT_URL} ...")
+    subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
+    subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
+
+    print("[INFO] Installing CUDA Toolkit silently ...")
+    subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])
+
+    print("[INFO] Setting CUDA environment variables ...")
+    os.environ["CUDA_HOME"] = "/usr/local/cuda"
+    os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ.get("PATH", ""))
+    os.environ["LD_LIBRARY_PATH"] = "%s/lib64:%s" % (
+        os.environ["CUDA_HOME"],
+        os.environ.get("LD_LIBRARY_PATH", "")
+    )
+
+    # Optional: set architecture list for compilation (Ampere and Ada)
+    os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6;8.9"
+    if os.path.exists(CUDA_TOOLKIT_FILE):
+        os.remove(CUDA_TOOLKIT_FILE)
+        print(f"[INFO] Removed installer file: {CUDA_TOOLKIT_FILE}")
+    else:
+        print(f"[WARN] Installer file not found: {CUDA_TOOLKIT_FILE}")
+    print(os.listdir("/usr/local/cuda"))
+    print("[INFO] CUDA 12.1 installation complete. CUDA_HOME set to /usr/local/cuda")
+logging.info("Environment Variables: %s" % os.environ)
+logging.info("Installing CUDA extensions...")
+if _get_output(["nvcc", "--version"]) is None:
+    logging.info("Installing CUDA toolkit...")
+    install_cuda_toolkit()
+    logging.info("installCUDA: %s" % _get_output(["nvcc", "--version"]))
+else:
+    logging.info("Detected CUDA: %s" % _get_output(["nvcc", "--version"]))
+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
@@ -56,47 +108,14 @@ import shutil
 
 # Suppress warnings (especially for PyTorch)
 warnings.filterwarnings("ignore")
-import os
-import subprocess
-import sys
 
-os.environ["MEDIAPIPE_DISABLE_GPU"] = "1"  # Disable GPU for MediaPipe
-def install_cuda_toolkit():
-    CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run"
-    CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
 
-    print(f"[INFO] Downloading CUDA Toolkit from {CUDA_TOOLKIT_URL} ...")
-    subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
-    subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
-
-    print("[INFO] Installing CUDA Toolkit silently ...")
-    subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])
-
-    print("[INFO] Setting CUDA environment variables ...")
-    os.environ["CUDA_HOME"] = "/usr/local/cuda"
-    os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ.get("PATH", ""))
-    os.environ["LD_LIBRARY_PATH"] = "%s/lib64:%s" % (
-        os.environ["CUDA_HOME"],
-        os.environ.get("LD_LIBRARY_PATH", "")
-    )
+os.environ["MEDIAPIPE_DISABLE_GPU"] = "1"  # Disable GPU for MediaPipe
 
-    # Optional: set architecture list for compilation (Ampere and Ada)
-    os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6;8.9"
-    if os.path.exists(CUDA_TOOLKIT_FILE):
-        os.remove(CUDA_TOOLKIT_FILE)
-        print(f"[INFO] Removed installer file: {CUDA_TOOLKIT_FILE}")
-    else:
-        print(f"[WARN] Installer file not found: {CUDA_TOOLKIT_FILE}")
-    print(os.listdir("/usr/local/cuda"))
-    print("[INFO] CUDA 12.1 installation complete. CUDA_HOME set to /usr/local/cuda")
 
 # 🔧 Set CUDA_HOME before anything else
 
-# Configure logging settings
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(levelname)s - %(message)s"
-)
+
 from diffusers import (
     StableDiffusionControlNetImg2ImgPipeline,
     ControlNetModel,
@@ -946,33 +965,25 @@ def launch_gradio_app():
 
         demo.queue()
         demo.launch(server_name="0.0.0.0")
-def _get_output(cmd):
-    try:
-        return subprocess.check_output(cmd).decode("utf-8")
-    except Exception as ex:
-        logging.exception(ex)
 
-    return None
 
 if __name__ == '__main__':
     import torch.multiprocessing as mp
     mp.set_start_method('spawn', force=True)
-    logging.info("Environment Variables: %s" % os.environ)
-    logging.info("Installing CUDA extensions...")
-    if _get_output(["nvcc", "--version"]) is None:
-        logging.info("Installing CUDA toolkit...")
-        install_cuda_toolkit()
-        logging.info("installCUDA: %s" % _get_output(["nvcc", "--version"]))
-
-    else:
-        logging.info("Detected CUDA: %s" % _get_output(["nvcc", "--version"]))
-
-    print("CUDA_HOME =", os.environ.get("CUDA_HOME"))
-    from torch.utils.cpp_extension import CUDA_HOME
-    print("CUDA_HOME from PyTorch:", CUDA_HOME)
+    # logging.info("Environment Variables: %s" % os.environ)
+    # logging.info("Installing CUDA extensions...")
+    # if _get_output(["nvcc", "--version"]) is None:
+    #     logging.info("Installing CUDA toolkit...")
+    #     install_cuda_toolkit()
+    #     logging.info("installCUDA: %s" % _get_output(["nvcc", "--version"]))
+
+    # else:
+    #     logging.info("Detected CUDA: %s" % _get_output(["nvcc", "--version"]))
+
+    # print("CUDA_HOME =", os.environ.get("CUDA_HOME"))
+    # from torch.utils.cpp_extension import CUDA_HOME
+    # print("CUDA_HOME from PyTorch:", CUDA_HOME)
     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, \