File size: 4,964 Bytes
7ee4d50 |
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 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
import os
from typing import List
import torch
from cog import BasePredictor, Input, Path
from diffusers import (
StableDiffusionPipeline,
PNDMScheduler,
LMSDiscreteScheduler,
DDIMScheduler,
EulerDiscreteScheduler,
EulerAncestralDiscreteScheduler,
DPMSolverMultistepScheduler,
)
from diffusers.pipelines.stable_diffusion.safety_checker import (
StableDiffusionSafetyChecker,
)
# MODEL_ID refers to a diffusers-compatible model on HuggingFace
# e.g. prompthero/openjourney-v2, wavymulder/Analog-Diffusion, etc
MODEL_ID = "stabilityai/stable-diffusion-2-1"
MODEL_CACHE = "diffusers-cache"
SAFETY_MODEL_ID = "CompVis/stable-diffusion-safety-checker"
class Predictor(BasePredictor):
def setup(self):
"""Load the model into memory to make running multiple predictions efficient"""
print("Loading pipeline...")
safety_checker = StableDiffusionSafetyChecker.from_pretrained(
SAFETY_MODEL_ID,
cache_dir=MODEL_CACHE,
local_files_only=True,
)
self.pipe = StableDiffusionPipeline.from_pretrained(
MODEL_ID,
safety_checker=safety_checker,
cache_dir=MODEL_CACHE,
local_files_only=True,
).to("cuda")
@torch.inference_mode()
def predict(
self,
prompt: str = Input(
description="Input prompt",
default="a photo of an astronaut riding a horse on mars",
),
negative_prompt: str = Input(
description="Specify things to not see in the output",
default=None,
),
width: int = Input(
description="Width of output image. Maximum size is 1024x768 or 768x1024 because of memory limits",
choices=[128, 256, 384, 448, 512, 576, 640, 704, 768, 832, 896, 960, 1024],
default=768,
),
height: int = Input(
description="Height of output image. Maximum size is 1024x768 or 768x1024 because of memory limits",
choices=[128, 256, 384, 448, 512, 576, 640, 704, 768, 832, 896, 960, 1024],
default=768,
),
num_outputs: int = Input(
description="Number of images to output.",
ge=1,
le=4,
default=1,
),
num_inference_steps: int = Input(
description="Number of denoising steps", ge=1, le=500, default=50
),
guidance_scale: float = Input(
description="Scale for classifier-free guidance", ge=1, le=20, default=7.5
),
scheduler: str = Input(
default="DPMSolverMultistep",
choices=[
"DDIM",
"K_EULER",
"DPMSolverMultistep",
"K_EULER_ANCESTRAL",
"PNDM",
"KLMS",
],
description="Choose a scheduler.",
),
seed: int = Input(
description="Random seed. Leave blank to randomize the seed", default=None
),
) -> List[Path]:
"""Run a single prediction on the model"""
if seed is None:
seed = int.from_bytes(os.urandom(2), "big")
print(f"Using seed: {seed}")
if width * height > 786432:
raise ValueError(
"Maximum size is 1024x768 or 768x1024 pixels, because of memory limits. Please select a lower width or height."
)
self.pipe.scheduler = make_scheduler(scheduler, self.pipe.scheduler.config)
generator = torch.Generator("cuda").manual_seed(seed)
output = self.pipe(
prompt=[prompt] * num_outputs if prompt is not None else None,
negative_prompt=[negative_prompt] * num_outputs
if negative_prompt is not None
else None,
width=width,
height=height,
guidance_scale=guidance_scale,
generator=generator,
num_inference_steps=num_inference_steps,
)
output_paths = []
for i, sample in enumerate(output.images):
if output.nsfw_content_detected and output.nsfw_content_detected[i]:
continue
output_path = f"/tmp/out-{i}.png"
sample.save(output_path)
output_paths.append(Path(output_path))
if len(output_paths) == 0:
raise Exception(
f"NSFW content detected. Try running it again, or try a different prompt."
)
return output_paths
def make_scheduler(name, config):
return {
"PNDM": PNDMScheduler.from_config(config),
"KLMS": LMSDiscreteScheduler.from_config(config),
"DDIM": DDIMScheduler.from_config(config),
"K_EULER": EulerDiscreteScheduler.from_config(config),
"K_EULER_ANCESTRAL": EulerAncestralDiscreteScheduler.from_config(config),
"DPMSolverMultistep": DPMSolverMultistepScheduler.from_config(config),
}[name]
|