Spaces:
Sleeping
Sleeping
File size: 3,761 Bytes
2978bba | 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 | #!/usr/bin/env python3
"""Helper script to download Stable Diffusion and LDM models for MorphGuard"""
import os
import argparse
import torch
from diffusers import StableDiffusionImg2ImgPipeline, DiffusionPipeline
def main():
parser = argparse.ArgumentParser(description="Download Stable Diffusion models for MorphGuard")
parser.add_argument("--use-cpu", action="store_true", help="Force CPU usage (not recommended)")
parser.add_argument("--no-safety-checker", action="store_true", help="Disable safety checker for faster operation")
parser.add_argument("--no-fp16", action="store_true", help="Disable fp16 even when running on GPU")
args = parser.parse_args()
# Determine device
if args.use_cpu:
device = "cpu"
print("WARNING: Forcing CPU usage. This will be very slow for inference.")
use_fp16 = False
else:
device = "cuda" if torch.cuda.is_available() else "cpu"
use_fp16 = device == "cuda" and not args.no_fp16
if device == "cpu":
print("WARNING: CUDA not available. Downloads will be configured for CPU but inference will be slow.")
else:
if use_fp16:
print(f"Using {device} with fp16 precision for model downloads (faster, lower memory usage)")
else:
print(f"Using {device} with fp32 precision (higher accuracy but more memory usage)")
# Models to download
models = [
{
"name": "Stable Diffusion v1.4",
"path": "models/diffusion/stable-diffusion-v1-4",
"model_id": "CompVis/stable-diffusion-v1-4",
"pipeline": StableDiffusionImg2ImgPipeline
},
{
"name": "LDM CelebA-HQ",
"path": "models/latent-diffusion/ldm-celebahq-256",
"model_id": "CompVis/ldm-celebahq-256",
"pipeline": DiffusionPipeline
}
]
# Ensure directories exist
for model in models:
os.makedirs(model["path"], exist_ok=True)
# Download each model
for model in models:
print(f"Downloading {model['name']} to {model['path']}...")
# Configure pipeline based on device and args
kwargs = {}
if args.no_safety_checker:
kwargs["safety_checker"] = None
if use_fp16:
# Use fp16 on CUDA if available (faster, lower memory usage)
kwargs["torch_dtype"] = torch.float16
kwargs["variant"] = "fp16"
print(f" - Using fp16 precision for {model['name']}")
# Create pipeline
try:
pipeline = model["pipeline"].from_pretrained(
model["model_id"],
**kwargs
)
# Move to device
pipeline = pipeline.to(device)
# Save to local path
pipeline.save_pretrained(model["path"])
print(f"✅ Successfully downloaded {model['name']}")
except Exception as e:
print(f"❌ Error downloading {model['name']}: {str(e)}")
print("\nSetup complete! Models were downloaded to:")
for model in models:
print(f" - {model['path']}")
if device == "cuda":
print("\nYour GPU has been detected and models configured for optimal performance.")
if use_fp16:
print("Using fp16 precision which offers:")
print(" - Up to 2-3x faster inference")
print(" - Half the memory usage")
print(" - Ability to run larger models or batch sizes")
else:
print("\nNo GPU detected. Consider running on a machine with CUDA support for better performance.")
if __name__ == "__main__":
main() |