File size: 2,661 Bytes
3324de2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
import sys
import os
import logging
import subprocess
import traceback
import json 
import re

log = logging.getLogger("TangoFlux")

download_models = True

EXT_PATH = os.path.dirname(os.path.abspath(__file__))

try:
    folder_paths_path = os.path.abspath(os.path.join(EXT_PATH, "..", "..", "..", "folder_paths.py"))

    sys.path.append(os.path.dirname(folder_paths_path))

    import folder_paths
    
    TANGOFLUX_DIR = os.path.join(folder_paths.models_dir, "tangoflux")
    TEXT_ENCODER_DIR = os.path.join(folder_paths.models_dir, "text_encoders")
except:
    download_models = False
    
try:
    log.info("Installing requirements")
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", f"{EXT_PATH}/requirements.txt", "--no-warn-script-location"])
    
    if download_models:
        from huggingface_hub import snapshot_download
        
        log.info("Downloading Necessary models")

        try:
            log.info(f"Downloading TangoFlux models to: {TANGOFLUX_DIR}")
            snapshot_download(
                repo_id="declare-lab/TangoFlux",
                allow_patterns=["*.json", "*.safetensors"],
                local_dir=TANGOFLUX_DIR,
                local_dir_use_symlinks=False,
            )
        except Exception:
            traceback.print_exc()
            log.error("Failed to download TangoFlux models")
            
        log.info("Loading config")

        with open(os.path.join(TANGOFLUX_DIR, "config.json"), "r") as f:
            config = json.load(f)
            
        try:
            text_encoder = re.sub(r'[<>:"/\\|?*]', '-', config.get("text_encoder_name", "google/flan-t5-large"))
            text_encoder_path = os.path.join(TEXT_ENCODER_DIR, text_encoder)
            
            log.info(f"Downloading text encoders to: {text_encoder_path}")
            snapshot_download(
                repo_id=config.get("text_encoder_name", "google/flan-t5-large"),
                allow_patterns=["*.json", "*.safetensors", "*.model"],
                local_dir=text_encoder_path,
                local_dir_use_symlinks=False,
            )
        except Exception:
            traceback.print_exc()
            log.error("Failed to download text encoders")
        
    try:
        log.info("Installing TangoFlux module")
        subprocess.check_call([sys.executable, "-m", "pip", "install", os.path.join(EXT_PATH, "..")])
    except Exception:
        traceback.print_exc()
        log.error("Failed to install TangoFlux module")
    
    log.info("TangoFlux Installation completed")
        
except Exception:
    traceback.print_exc()
    log.error("TangoFlux Installation failed")