Spaces:
Runtime error
Runtime error
File size: 1,753 Bytes
d51d951 2df9553 96d24c8 2df9553 4b90f4d 3206fc4 96d24c8 4b90f4d 2df9553 627164f 5126a09 627164f d7623eb 2df9553 1e0e2d1 1ecf68a 2df9553 96d24c8 2df9553 96d24c8 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 |
from __future__ import annotations
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
from diffusers import DPMSolverMultistepScheduler
import torch
import PIL.Image
import numpy as np
# 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 = "runwayml/stable-diffusion-v1-5"
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 = 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.int64).max)
generator = torch.Generator(device).manual_seed(seed)
return 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
# image = pipeline(prompt=prompt,
# negative_prompt = n_prompt,
# num_inference_steps = 2,
# guidance_scale = 7).images
|