RetratoWaifuStudio / utils.py
Josue Aaron Soriano Rivero
fix: removido anime screenshot reemplazado por cute
a0541f7
import torch
from torch.cuda.amp import autocast
from translations import CATEGORIES, TRANSLATIONS
def build_prompt(**kwargs):
default_params = ["masterpiece", "best quality", "1girl", "watercolor", "portrait", "upper body", "looking at viewer", "cute", "highres"]
for key, values in kwargs.items():
if values:
english_values = [CATEGORIES[key][TRANSLATIONS[key].index(v)] for v in values if v in TRANSLATIONS[key]]
default_params.append(", ".join(english_values))
return ', '.join(default_params)
def generate_image(prompt, pipe, device):
height = 512
width = 512
guidance_scale = 15
num_inference_steps = 50
pipe.enable_attention_slicing()
with torch.autocast('cuda') if device == 'cuda' else torch.inference_mode():
image = pipe(
prompt,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
height=height,
width=width
).images[0]
return image