Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,515 Bytes
6aefd85 781a759 6aefd85 781a759 b82c0ac 23063b3 58cfe29 7cd4941 |
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 |
import gc
import os
import random
import numpy as np
import json
import torch
import uuid
from PIL import Image, PngImagePlugin
from datetime import datetime
from dataclasses import dataclass
from typing import Callable, Dict, Optional, Tuple, Any, List
from diffusers import (
DDIMScheduler,
DPMSolverMultistepScheduler,
DPMSolverSinglestepScheduler,
EulerAncestralDiscreteScheduler,
EulerDiscreteScheduler,
AutoencoderKL,
StableDiffusionXLPipeline,
)
import logging
def load_pipeline(model_name: str, device: torch.device, hf_token: Optional[str] = None, vae: Optional[AutoencoderKL] = None) -> Any:
"""Load the Stable Diffusion pipeline."""
try:
pipeline = (
StableDiffusionXLPipeline.from_single_file
if model_name.endswith(".safetensors")
else StableDiffusionXLPipeline.from_pretrained
)
pipe = pipeline(
model_name,
vae=vae,
torch_dtype=torch.float16,
custom_pipeline="lpw_stable_diffusion_xl",
use_safetensors=True,
add_watermarker=False
)
pipe.to(device)
return pipe
except Exception as e:
logging.error(f"Failed to load pipeline: {str(e)}", exc_info=True)
raise
def seed_everything(seed: int) -> torch.Generator:
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
generator = torch.Generator()
generator.manual_seed(seed)
return generator
def preprocess_image_dimensions(width, height):
if width % 8 != 0:
width = width - (width % 8)
if height % 8 != 0:
height = height - (height % 8)
return width, height
def get_scheduler(scheduler_config: Dict, name: str) -> Optional[Callable]:
scheduler_factory_map = {
"DPM++ 2M Karras": lambda: DPMSolverMultistepScheduler.from_config(
scheduler_config, use_karras_sigmas=True
),
"DPM++ SDE Karras": lambda: DPMSolverSinglestepScheduler.from_config(
scheduler_config, use_karras_sigmas=True
),
"DPM++ 2M SDE Karras": lambda: DPMSolverMultistepScheduler.from_config(
scheduler_config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++"
),
"Euler": lambda: EulerDiscreteScheduler.from_config(scheduler_config),
"Euler a": lambda: EulerAncestralDiscreteScheduler.from_config(
scheduler_config
),
"DDIM": lambda: DDIMScheduler.from_config(scheduler_config),
}
return scheduler_factory_map.get(name, lambda: None)()
def common_upscale(
samples: torch.Tensor,
width: int,
height: int,
upscale_method: str,
) -> torch.Tensor:
return torch.nn.functional.interpolate(
samples, size=(height, width), mode=upscale_method
)
def upscale(
samples: torch.Tensor, upscale_method: str, scale_by: float
) -> torch.Tensor:
width = round(samples.shape[3] * scale_by)
height = round(samples.shape[2] * scale_by)
return common_upscale(samples, width, height, upscale_method)
def free_memory() -> None:
"""Free up GPU and system memory."""
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
gc.collect()
def save_image(image, output_dir):
filename = str(uuid.uuid4()) + ".jpg"
os.makedirs(output_dir, exist_ok=True)
filepath = os.path.join(output_dir, filename)
image.save(filepath, "JPEG", quality=80)
return filepath |