File size: 2,077 Bytes
14d83f6 |
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 |
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()}")
# True
print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
# Tesla T4
device = "cuda"
class Model:
def __init__(self, modelID):
#modelID = "runwayml/stable-diffusion-v1-5"
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()
#self.pipe = StableDiffusionPipeline.from_pretrained(modelID)
#prompt = "a photo of an astronaut riding a horse on mars"
#n_prompt = "deformed, disfigured"
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
# image = pipeline(prompt=prompt,
# negative_prompt = n_prompt,
# num_inference_steps = 2,
# guidance_scale = 7).images
|