Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,348 Bytes
b2682d8 76cfa14 b2682d8 |
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 |
import os
import torch
from huggingface_hub import snapshot_download
from diffusers import StableDiffusionBrushNetPipeline, BrushNetModel, UniPCMultistepScheduler
torch_dtype = torch.float16
device = "cpu"
BrushEdit_path = "models/"
if not os.path.exists(BrushEdit_path):
BrushEdit_path = snapshot_download(
repo_id="TencentARC/BrushEdit",
local_dir=BrushEdit_path,
token=os.getenv("HF_TOKEN"),
)
brushnet_path = os.path.join(BrushEdit_path, "brushnetX")
brushnet = BrushNetModel.from_pretrained(brushnet_path, torch_dtype=torch_dtype)
base_models_list = [
# {
# "name": "dreamshaper_8 (Preload)",
# "local_path": "models/base_model/dreamshaper_8",
# "pipe": StableDiffusionBrushNetPipeline.from_pretrained(
# "models/base_model/dreamshaper_8", brushnet=brushnet, torch_dtype=torch_dtype, low_cpu_mem_usage=False
# ).to(device)
# },
# {
# "name": "epicrealism (Preload)",
# "local_path": "models/base_model/epicrealism_naturalSinRC1VAE",
# "pipe": StableDiffusionBrushNetPipeline.from_pretrained(
# "models/base_model/epicrealism_naturalSinRC1VAE", brushnet=brushnet, torch_dtype=torch_dtype, low_cpu_mem_usage=False
# ).to(device)
# },
{
"name": "henmixReal (Preload)",
"local_path": "models/base_model/henmixReal_v5c",
"pipe": StableDiffusionBrushNetPipeline.from_pretrained(
"models/base_model/henmixReal_v5c", brushnet=brushnet, torch_dtype=torch_dtype, low_cpu_mem_usage=False
).to(device)
},
{
"name": "meinamix (Preload)",
"local_path": "models/base_model/meinamix_meinaV11",
"pipe": StableDiffusionBrushNetPipeline.from_pretrained(
"models/base_model/meinamix_meinaV11", brushnet=brushnet, torch_dtype=torch_dtype, low_cpu_mem_usage=False
).to(device)
},
{
"name": "realisticVision (Default)",
"local_path": "models/base_model/realisticVisionV60B1_v51VAE",
"pipe": StableDiffusionBrushNetPipeline.from_pretrained(
"models/base_model/realisticVisionV60B1_v51VAE", brushnet=brushnet, torch_dtype=torch_dtype, low_cpu_mem_usage=False
).to(device)
},
]
base_models_template = {k["name"]: (k["local_path"], k["pipe"]) for k in base_models_list}
|