File size: 2,295 Bytes
fc9acd0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
from factories import UNet_conditional
from wrapper import DiffusionManager, Schedule
import os
import re
import torch
from bert_vectorize import vectorize_text_with_bert, cleanup
import time
import torchvision
from logger import save_grid_with_label
from clip_score import select_top_n_images
from torchinfo import summary
EXPERIMENT_DIRECTORY = "runs/run_3_jxa_resumed"
device = "mps" if torch.backends.mps.is_available() else "cpu"
try:
os.mkdir(os.path.join(EXPERIMENT_DIRECTORY, "inferred"))
except:
print("Skipping making directory, directory already exists")
net = UNet_conditional(num_classes=768)
net.to(device)
net.load_state_dict(torch.load(os.path.join(EXPERIMENT_DIRECTORY, "ckpt/latest.pt")))
def count_parameters(model):
return torch.tensor([p.numel() for p in model.parameters() if p.requires_grad]).sum().item()
print(f"Parameters: {count_parameters(net)}")
wrapper = DiffusionManager(net, device=device, noise_steps=1000)
wrapper.set_schedule(Schedule.LINEAR)
def infer(prompt, amt=1, topn=8):
path = os.path.join(EXPERIMENT_DIRECTORY, "inferred", re.sub(r'[^a-zA-Z\s]', '', prompt).replace(" ", "_")+str(int(time.time()))+".png")
vprompt = vectorize_text_with_bert(prompt).unsqueeze(0)
generated = wrapper.sample(64, vprompt, amt=amt).detach().cpu()
generated, _ = select_top_n_images(generated, prompt, n=topn)
save_grid_with_label(torchvision.utils.make_grid(generated),prompt + f"({topn} best of {amt})", path)
def run_jobs():
n=8
bestof=32
print(f"using best {bestof} of {n}")
processed_tasks = set()
def read_jobs():
try:
with open("inference_jobs.txt", 'r') as file:
tasks = file.readlines()
return [task.strip() for task in tasks]
except FileNotFoundError:
return []
tasks = read_jobs()
new_tasks = [task for task in tasks if task not in processed_tasks]
while new_tasks:
if new_tasks:
for task in new_tasks:
infer(task, n,bestof)
processed_tasks.add(task)
tasks = read_jobs()
new_tasks = [task for task in tasks if task not in processed_tasks]
cleanup()
if __name__ == "__main__":
#infer(input("Prompt? "), 8)
run_jobs() |