22

app.py

@@ -8,10 +8,15 @@ from transformers import AutoTokenizer, Qwen3ForCausalLM
 from diffusers import AutoencoderKL
 from utils.image_utils import get_image_latent, rescale_image
 from utils.prompt_utils import polish_prompt
+from utils import repo_utils
 # from controlnet_aux import HEDdetector, MLSDdetector, OpenposeDetector, CannyDetector, MidasDetector
 from controlnet_aux.processor import Processor
 
 
+repo_utils.clone_repo_if_not_exists("https://huggingface.co/Tongyi-MAI/Z-Image-Turbo", "app/models")
+repo_utils.clone_repo_if_not_exists("https://huggingface.co/alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union", "app/models")
+
+
 # MODEL_REPO = "Tongyi-MAI/Z-Image-Turbo"
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1280

utils/repo_utils.py

@@ -0,0 +1,33 @@
+import os, subprocess
+
+def clone_repo_if_not_exists(git_url, git_dir):
+    """
+    Clones a Git repository if it doesn't already exist in the git_dir folder.
+    """
+    home_dir = os.path.expanduser("~")
+    models_dir = os.path.join(home_dir, git_dir)
+
+    # Extract repository name from the Git URL
+    if git_url.endswith(".git"):
+        git_name = git_url.split('/')[-1][:-4]
+    else:
+        git_name = git_url.split('/')[-1]
+
+    repo_path = os.path.join(models_dir, git_name)
+
+    if not os.path.exists(models_dir):
+        os.makedirs(models_dir)
+        print(f"Created directory: {models_dir}")
+
+    if not os.path.exists(repo_path):
+        print(f"Repository '{git_name}' not found in '{models_dir}'. Cloning from {git_url}...")
+        try:
+            subprocess.run(["git", "clone", git_url, repo_path], check=True)
+            print(f"Successfully cloned '{git_name}' to '{repo_path}'.")
+        except subprocess.CalledProcessError as e:
+            print(f"Error cloning repository: {e}")
+        except FileNotFoundError:
+            print("Error: 'git' command not found. Please ensure Git is installed and in your PATH.")
+    else:
+        print(f"Repository '{git_name}' already exists at '{repo_path}'. Skipping clone.")
+