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