import torch
import time
import os
import pandas as pd
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler, EulerDiscreteScheduler, DDIMScheduler, DiffusionPipeline

model_id = "stable-diffusion-2-1-base/"
text_file = "PartiPrompts.tsv"

df = pd.read_csv(text_file, sep='\t')
prompts = df['Prompt']
num_inference_steps = 50
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")
base_count = 0
print("======================================start DPM ==================================")
for prompt in prompts:
    start = time.time()
    image = pipe(prompt, 512, 512, num_inference_steps=num_inference_steps, num_images_per_prompt=1).images[0]
    print(f"the {base_count} text-to-image use time {time.time()-start}")
    base_count += 1
    image.save(f"{base_count:05}.png")
    if base_count == 1:
        break
    
print(f"Your samples are ready and waiting for you here\n \n"
          f" \nEnjoy.")