from __future__ import annotations


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

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