Spaces:
Running
Running
| #!/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() |