from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
import torch
import os


def gen_img_from_text(prompts, gen_class, save_dir, max_num=5):

    repo_id = "stabilityai/stable-diffusion-2-base"
    pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.float16, variant="fp16")

    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    pipe = pipe.to("cuda")

    for prompt in prompts:
        images = pipe(prompt, num_inference_steps=50, num_images_per_prompt=max_num).images
        save_path = os.path.join(save_dir, f"{gen_class}", prompt.replace(". Specifically, it ", "").replace(",", "").replace(" ", "_").replace("This picture associates with ", ""))
        os.makedirs(save_path, exist_ok=True)
        for iim, image in enumerate(images):
            image.save(os.path.join(save_path, f"00000{iim}.jpg"))
