Spaces:
Paused
Paused
| from __future__ import annotations | |
| from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler | |
| from diffusers import DPMSolverMultistepScheduler | |
| import torch | |
| import PIL.Image | |
| import numpy as np | |
| import datetime | |
| # Check environment | |
| print(f"Is CUDA available: {torch.cuda.is_available()}") | |
| print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}") | |
| device = "cuda" | |
| class Model: | |
| def __init__(self, modelID): | |
| self.modelID = modelID | |
| self.pipe = StableDiffusionPipeline.from_pretrained(modelID, torch_dtype=torch.float16) | |
| self.pipe = self.pipe.to(device) | |
| self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config) | |
| self.pipe.enable_xformers_memory_efficient_attention() | |
| def process(self, | |
| prompt: str, | |
| negative_prompt: str, | |
| guidance_scale:int = 7, | |
| num_images:int = 1, | |
| num_steps:int = 20, | |
| ): | |
| seed = np.random.randint(0, np.iinfo(np.int32).max) | |
| generator = torch.Generator(device).manual_seed(seed) | |
| now = datetime.datetime.now() | |
| print(now) | |
| print(self.modelID) | |
| print(prompt) | |
| print(negative_prompt) | |
| with torch.inference_mode(): | |
| images = self.pipe(prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| guidance_scale=guidance_scale, | |
| num_images_per_prompt=num_images, | |
| num_inference_steps=num_steps, | |
| generator=generator).images | |
| return images | |