ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-29 10:41:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixed Dora implementation. Still highly experimental

Browse Source
This commit is contained in:
Jaret Burkett
2024-02-24 10:26:01 -07:00
parent 1bd94f0f01
commit f965a1299f
5 changed files with 128 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
jobs/process/GenerateProcess.py
View File

@@ -32,16 +32,30 @@ class GenerateConfig:
self.guidance_rescale = kwargs.get('guidance_rescale', 0.0)
self.ext = kwargs.get('ext', 'png')
self.prompt_file = kwargs.get('prompt_file', False)
self.prompts_in_file = self.prompts
if self.prompts is None:
raise ValueError("Prompts must be set")
if isinstance(self.prompts, str):
if os.path.exists(self.prompts):
with open(self.prompts, 'r', encoding='utf-8') as f:
self.prompts = f.read().splitlines()
self.prompts = [p.strip() for p in self.prompts if len(p.strip()) > 0]
self.prompts_in_file = f.read().splitlines()
self.prompts_in_file = [p.strip() for p in self.prompts_in_file if len(p.strip()) > 0]
else:
raise ValueError("Prompts file does not exist, put in list if you want to use a list of prompts")
self.random_prompts = kwargs.get('random_prompts', False)
self.max_random_per_prompt = kwargs.get('max_random_per_prompt', 1)
self.max_images = kwargs.get('max_prompts', 10000)
if self.random_prompts:
self.prompts = []
for i in range(self.max_images):
num_prompts = random.randint(1, self.max_random_per_prompt)
prompt_list = [random.choice(self.prompts_in_file) for _ in range(num_prompts)]
self.prompts.append(", ".join(prompt_list))
else:
self.prompts = self.prompts_in_file
if kwargs.get('shuffle', False):
# shuffle the prompts
random.shuffle(self.prompts)
@@ -78,6 +92,9 @@ class GenerateProcess(BaseProcess):
print("Loading model...")
self.sd.load_model()
print("Compiling model...")
self.sd.unet = torch.compile(self.sd.unet, mode="reduce-overhead", fullgraph=True)
print(f"Generating {len(self.generate_config.prompts)} images")
# build prompt image configs
prompt_image_configs = []