Spaces:
Sleeping
Sleeping
| import argparse | |
| import os | |
| import secrets | |
| import threading | |
| from pathlib import Path | |
| from PIL import Image | |
| ROOT = Path(__file__).resolve().parent | |
| MODEL_ID = "google/ddpm-celebahq-256" | |
| MODEL_SIZES = {MODEL_ID: 256, "google/ddpm-cifar10-32": 32} | |
| CHECKPOINTS = (0, 50, 100, 200, 350, 500, 700, 1000) | |
| # Gradio may resume a generator on a different worker thread between yields. | |
| # A primitive Lock can be released there; RLock requires the acquiring thread. | |
| _MODEL_LOCK = threading.Lock() | |
| _CACHE_LOCK = threading.Lock() | |
| _PIPELINES = {} | |
| _PIPELINE = None | |
| _TORCH = None | |
| _DEVICE = None | |
| def load_model(device=None, model_id=MODEL_ID): | |
| """Load and cache the diffusion pipeline, returning it and its device.""" | |
| global _PIPELINE, _TORCH, _DEVICE | |
| if model_id not in MODEL_SIZES: | |
| raise ValueError(f"Unsupported model: {model_id}") | |
| with _CACHE_LOCK: | |
| try: | |
| import torch | |
| from diffusers import DDPMPipeline | |
| except ImportError as error: | |
| raise RuntimeError( | |
| "Diffusers dependencies are missing from this Python environment. " | |
| "Install them with `python -m pip install -r requirements.txt`. " | |
| f"Import detail: {error}" | |
| ) from error | |
| if device is None: | |
| if torch.cuda.is_available(): | |
| device = "cuda" | |
| elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): | |
| device = "mps" | |
| else: | |
| device = "cpu" | |
| device = str(device) | |
| key = (model_id, device) | |
| if key not in _PIPELINES: | |
| print(f"Loading {model_id} on device: {device}") | |
| pipeline = DDPMPipeline.from_pretrained( | |
| model_id, cache_dir=Path(os.environ.get("MODEL_CACHE_DIR", ROOT / "models"))) | |
| if device == "cuda" and os.environ.get("SPACES_ZERO_GPU", "").lower() in {"1", "true", "yes"}: | |
| # ZeroGPU packs CUDA aliases at startup. Give it owned CPU | |
| # storage rather than tensors backed by mapped checkpoint files. | |
| for component in pipeline.components.values(): | |
| if isinstance(component, torch.nn.Module): | |
| component._apply(lambda tensor: tensor.clone()) | |
| pipeline.to(device) | |
| _PIPELINES[key] = pipeline | |
| _PIPELINE = _PIPELINES[key] | |
| _TORCH = torch | |
| _DEVICE = device | |
| return _PIPELINE, _TORCH, _DEVICE | |
| def is_model_loaded(): | |
| return _PIPELINE is not None | |
| def sample_to_image(sample): | |
| """Convert a model sample tensor into a PIL image.""" | |
| image = sample[0].detach().float().cpu().clamp(-1, 1) | |
| image = ((image + 1) * 127.5).round().byte().permute(1, 2, 0).numpy() | |
| return Image.fromarray(image, mode="RGB") | |
| def generate_images(seed, checkpoints=CHECKPOINTS, initial_image=None, device=None, model_id=MODEL_ID): | |
| """Yield every remaining-step count and an image at requested checkpoints.""" | |
| with _MODEL_LOCK: | |
| pipeline, torch, device = load_model(device=device, model_id=model_id) | |
| scheduler = pipeline.scheduler | |
| training_steps = scheduler.config.num_train_timesteps | |
| scheduler.set_timesteps(training_steps, device=device) | |
| # MPS has no torch.Generator implementation. Keep its random stream | |
| # on CPU; Diffusers also accepts this generator for scheduler noise. | |
| random_device = "cpu" if device == "mps" else device | |
| generator = torch.Generator(device=random_device).manual_seed(seed) | |
| unet = pipeline.unet | |
| sample_size = unet.config.sample_size | |
| if isinstance(sample_size, (tuple, list)): | |
| height, width = sample_size | |
| else: | |
| height = width = sample_size | |
| if initial_image is None: | |
| sample = torch.randn( | |
| (1, unet.config.in_channels, height, width), | |
| generator=generator, | |
| device=random_device, | |
| dtype=unet.dtype, | |
| ).to(device) * scheduler.init_noise_sigma | |
| else: | |
| import numpy as np | |
| pixels = np.asarray(initial_image.convert("RGB").resize((width, height))).copy() | |
| sample = torch.from_numpy(pixels).permute(2, 0, 1).unsqueeze(0) | |
| sample = sample.to(device=device, dtype=unet.dtype) / 127.5 - 1 | |
| yield ( | |
| training_steps, | |
| sample_to_image(sample) if training_steps in checkpoints else None, | |
| ) | |
| with torch.inference_mode(): | |
| for index, timestep in enumerate(scheduler.timesteps): | |
| model_output = unet(sample, timestep).sample | |
| sample = scheduler.step( | |
| model_output, | |
| timestep, | |
| sample, | |
| generator=generator, | |
| ).prev_sample | |
| remaining_steps = training_steps - index - 1 | |
| image = sample_to_image(sample) if remaining_steps in checkpoints else None | |
| yield remaining_steps, image | |
| def generate_image(seed, model_id=MODEL_ID): | |
| """Generate and return the final image for a seed.""" | |
| final_image = None | |
| for _, image in generate_images(seed, checkpoints=(0,), model_id=model_id): | |
| if image is not None: | |
| final_image = image | |
| if final_image is None: | |
| raise RuntimeError("The diffusion pipeline did not produce an image.") | |
| return final_image | |
| def main(): | |
| parser = argparse.ArgumentParser(description="Generate an unconditional DDPM image.") | |
| parser.add_argument( | |
| "--seed", | |
| type=int, | |
| default=None, | |
| help="Random seed (defaults to a randomly generated seed)", | |
| ) | |
| parser.add_argument("--model", choices=list(MODEL_SIZES), default=MODEL_ID) | |
| args = parser.parse_args() | |
| seed = args.seed if args.seed is not None else secrets.randbelow(2**31) | |
| if seed < 0 or seed >= 2**63: | |
| parser.error("seed must be between 0 and 2^63 - 1") | |
| print(f"Generating with seed: {seed}") | |
| image = generate_image(seed, model_id=args.model) | |
| import matplotlib.pyplot as plt | |
| plt.imshow(image) | |
| plt.title(f"{args.model} — seed {seed}") | |
| plt.axis("off") | |
| plt.tight_layout() | |
| plt.show() | |
| if __name__ == "__main__": | |
| main() | |