import torch
from diffusers import StableDiffusionPipeline, StableDiffusion3Pipeline
import time

start_timestamp = time.time()
#model = "stabilityai/stable-diffusion-3.5-medium" # generation time: 13 min
model = "stabilityai/stable-diffusion-3-medium-diffusers" # generation time: 10 min
#model = "stabilityai/stable-diffusion-2" # generation time: 4 sec

pipe = StableDiffusion3Pipeline.from_pretrained(model, torch_dtype=torch.float16)
#pipe = StableDiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

model_loaded_timestamp = time.time()
model_load_time = model_loaded_timestamp - start_timestamp
print(f"model load time: {round(model_load_time)} seconds")

prompt = "A majestic broadsword with a golden pommel, no background"
image = pipe(
    prompt,
    guidance_scale=3.0,
).images[0]

image_name = "image7.png"
image.save(f"images/{image_name}")

generation_time = time.time() - model_loaded_timestamp
print(f"image generation time: {round(generation_time)} seconds")