bala1802/StableDiffusionModel

1import torch2import gc3 4import utils5import model6import config7import image_generator as generator8 9def predict(prompt, pipe, loss_function=None):10    latents = []11    12    for seed_number, sd_concept in zip(config.SEEDS, config.STABLE_DIFUSION_CONCEPTS):13        # torch.cuda.empty_cache()14        gc.collect()15        # torch.cuda.empty_cache()16        17        prompt = [f'{prompt} {sd_concept}']18        latent = generator.generate_images(pipe=pipe, seed_number=seed_number, prompt=prompt, loss_function=loss_function)19        latents.append(latent)20    21    latents = torch.vstack(latents)22    images = utils.convert_latents_to_pil_images(pipe=pipe, latents=latents)23    grid = utils.populate_image_grid(images, 1, len(latents))24    return grid25