Update lollms_diffusers.py

This commit is contained in:
Saifeddine ALOUI 2024-09-27 16:25:34 +02:00 committed by GitHub
parent 242bdb68dc
commit 97f237360b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -238,13 +238,31 @@ class LollmsDiffusers(LollmsTTI):
self.model.scheduler = sc self.model.scheduler = sc
width = adjust_dimensions(int(width)) width = adjust_dimensions(int(width))
height = adjust_dimensions(int(height)) height = adjust_dimensions(int(height))
def process_output_path(output_path, self_output_dir):
if output_path is None:
output_path = Path(self_output_dir)
fn = find_next_available_filename(output_path, "diff_img_")
else:
output_path = Path(output_path)
if output_path.is_file():
fn = output_path
elif output_path.is_dir():
fn = find_next_available_filename(output_path, "diff_img_")
else:
# If the path doesn't exist, assume it's intended to be a file
fn = output_path
return fn
# Usage in the original context
if output_path is None: if output_path is None:
output_path = self.output_dir output_path = self.output_dir
output_path = Path(output_path)
fn = find_next_available_filename(output_path,"diff_img_") fn = process_output_path(output_path, self.output_dir)
else:
output_path = Path(output_path)
fn = output_path
if seed!=-1: if seed!=-1:
generator = torch.Generator("cuda").manual_seed(seed) generator = torch.Generator("cuda").manual_seed(seed)
image = self.model(positive_prompt, negative_prompt=negative_prompt, height=height, width=width, guidance_scale=scale, num_inference_steps=steps, generator=generator).images[0] image = self.model(positive_prompt, negative_prompt=negative_prompt, height=height, width=width, guidance_scale=scale, num_inference_steps=steps, generator=generator).images[0]