Spaces:
Running
on
Zero
Running
on
Zero
File size: 6,728 Bytes
8b60c80 72be7ba ff5d480 72be7ba 5f435f3 72be7ba 8b60c80 |
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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 |
import os
from diffusers import DiffusionPipeline, AutoencoderTiny
import torch
# Define models and their configurations
models = {
"FLUX.1-dev": {
"pipeline_class": DiffusionPipeline,
"model_id": "black-forest-labs/FLUX.1-dev",
"config": {"torch_dtype": torch.bfloat16},
"description": "**FLUX.1-dev** is a development model that focuses on delivering highly detailed and artistically rich images.",
},
}
# Helper function to get the Hugging Face token securely
def get_hf_token():
try:
from google.colab import userdata # Try to get token from Colab secrets
hf_token = userdata.get('HF_TOKEN')
if hf_token:
return hf_token
else:
raise RuntimeError("HF_TOKEN not found in Colab secrets.")
except ImportError: # Not running in Colab
return os.getenv("HF_TOKEN", None)
# Function to pre-download models
def download_all_models():
print("Downloading all models...")
_HF_TOKEN = get_hf_token()
if not _HF_TOKEN:
raise ValueError("HF_TOKEN is not available. Please set it in Colab secrets or environment variables.")
for model_key, config in models.items():
try:
pipeline_class = config["pipeline_class"]
model_id = config["model_id"]
# Download the pipeline (weights will be cached)
pipeline_class.from_pretrained(model_id, token=_HF_TOKEN, **config.get("config", {}))
print(f"Model '{model_key}' downloaded successfully.")
except Exception as e:
print(f"Error downloading model '{model_key}': {e}")
print("Model download process complete.")
# Download the only VAE needed
print("Downloading VAE...")
try:
AutoencoderTiny.from_pretrained("madebyollin/taef1", token=_HF_TOKEN)
print("VAE 'taef1' downloaded successfully.")
except Exception as e:
print(f"Error downloading VAE: {e}")
print("VAE download process complete.")
'''
import os
from diffusers import DiffusionPipeline, AutoencoderTiny
import torch
# Define models and their configurations
models = {
"FLUX.1-dev": {
"pipeline_class": DiffusionPipeline,
"model_id": "black-forest-labs/FLUX.1-dev",
"config": {"torch_dtype": torch.bfloat16},
"description": "**FLUX.1-dev** is a development model that focuses on delivering highly detailed and artistically rich images.",
},
}
# Helper function to get the Hugging Face token securely
def get_hf_token():
try:
from google.colab import userdata # Try to get token from Colab secrets
hf_token = userdata.get('HF_TOKEN')
if hf_token:
return hf_token
else:
raise RuntimeError("HF_TOKEN not found in Colab secrets.")
except ImportError: # Not running in Colab
return os.getenv("HF_TOKEN", None)
# Function to pre-download models
def download_all_models():
print("Downloading all models...")
_HF_TOKEN = get_hf_token()
if not _HF_TOKEN:
raise ValueError("HF_TOKEN is not available. Please set it in Colab secrets or environment variables.")
for model_key, config in models.items():
try:
pipeline_class = config["pipeline_class"]
model_id = config["model_id"]
# Download the pipeline (weights will be cached)
pipeline_class.from_pretrained(model_id, token=_HF_TOKEN, **config.get("config", {}))
print(f"Model '{model_key}' downloaded successfully.")
except Exception as e:
print(f"Error downloading model '{model_key}': {e}")
print("Model download process complete.")
# Download the only VAE needed
print("Downloading VAE...")
try:
AutoencoderTiny.from_pretrained("madebyollin/taef1", token=_HF_TOKEN)
print("VAE 'taef1' downloaded successfully.")
except Exception as e:
print(f"Error downloading VAE: {e}")
print("VAE download process complete.")
import os
from diffusers import DiffusionPipeline, FluxPipeline, AutoencoderTiny, AutoencoderKL
from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
import torch
# Define models and their configurations (same as in app.py)
models = {
"FLUX.1-schnell": {
"pipeline_class": FluxPipeline,
"model_id": "black-forest-labs/FLUX.1-schnell",
"config": {"torch_dtype": torch.bfloat16},
"description": "**FLUX.1-schnell** is a fast and efficient model designed for quick image generation.",
},
"FLUX.1-dev": {
"pipeline_class": DiffusionPipeline,
"model_id": "black-forest-labs/FLUX.1-dev",
"config": {"torch_dtype": torch.bfloat16},
"description": "**FLUX.1-dev** is a development model that focuses on delivering highly detailed and artistically rich images.",
},
}
# Helper function to get the Hugging Face token securely
def get_hf_token():
try:
from google.colab import userdata # Try to get token from Colab secrets
hf_token = userdata.get('HF_TOKEN')
if hf_token:
return hf_token
else:
raise RuntimeError("HF_TOKEN not found in Colab secrets.")
except ImportError: # Not running in Colab
return os.getenv("HF_TOKEN", None)
# Function to pre-download models
def download_all_models():
print("Downloading all models...")
_HF_TOKEN = get_hf_token() # Get the token once
if not _HF_TOKEN:
raise ValueError("HF_TOKEN is not available. Please set it in Colab secrets or environment variables.")
for model_key, config in models.items():
try:
pipeline_class = config["pipeline_class"]
model_id = config["model_id"]
# Download the pipeline (weights will be cached)
pipeline_class.from_pretrained(model_id, token=_HF_TOKEN, **config.get("config", {}))
print(f"Model '{model_key}' downloaded successfully.")
except Exception as e:
print(f"Error downloading model '{model_key}': {e}")
print("Model download process complete.")
def download_vaes():
print("Downloading VAEs...")
try:
# Download taef1
AutoencoderTiny.from_pretrained("madebyollin/taef1", use_auth_token=get_hf_token())
print("VAE 'taef1' downloaded successfully.")
# Download good_vae (AutoencoderKL from FLUX.1-dev)
AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", use_auth_token=get_hf_token())
print("VAE 'good_vae' downloaded successfully.")
except Exception as e:
print(f"Error downloading VAEs: {e}")
print("VAE download process complete.")
''' |