File size: 1,555 Bytes
d51d951
 
 
2df9553
96d24c8
2df9553
 
 
 
96d24c8
 
2df9553
 
 
96d24c8
 
 
 
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
from __future__ import annotations


from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
from diffusers import DPMSolverMultistepScheduler
import torch
import PIL.Image
import numpy as np

device = "cpu"

class Model:
    def __init__(self):
        modelID = "runwayml/stable-diffusion-v1-5"
        pipe = StableDiffusionPipeline.from_pretrained(modelID, torch_dtype=torch.float16)
        pipe = pipe.to(device)
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(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