Spaces:
Runtime error
Runtime error
File size: 2,077 Bytes
d51d951 2df9553 96d24c8 2df9553 99671bd 2df9553 4b90f4d 3206fc4 96d24c8 4b90f4d 2df9553 99671bd 627164f 5126a09 99671bd 627164f d7623eb 2df9553 1e0e2d1 1ecf68a 99671bd 96d24c8 99671bd 5852deb 99671bd 2df9553 96d24c8 2df9553 99671bd 2df9553 |
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
|