add omost.py + omost_flux_example

This commit is contained in:
tc2000731
2024-09-03 19:40:40 +08:00
parent fe485b3fa1
commit 0b066d3cb4
7 changed files with 398 additions and 13 deletions

View File

@@ -1,8 +1,7 @@
from transformers import AutoTokenizer
from ..models.model_manager import ModelManager
import torch
from .omost import OmostPromter
class BeautifulPrompt(torch.nn.Module):
def __init__(self, tokenizer_path=None, model=None, template=""):
@@ -13,8 +12,8 @@ class BeautifulPrompt(torch.nn.Module):
@staticmethod
def from_model_manager(model_nameger: ModelManager):
model, model_path = model_nameger.fetch_model("beautiful_prompt", require_model_path=True)
def from_model_manager(model_manager: ModelManager):
model, model_path = model_manager.fetch_model("beautiful_prompt", require_model_path=True)
template = 'Instruction: Give a simple description of the image to generate a drawing prompt.\nInput: {raw_prompt}\nOutput:'
if model_path.endswith("v2"):
template = """Converts a simple image description into a prompt. \
@@ -63,8 +62,8 @@ class Translator(torch.nn.Module):
@staticmethod
def from_model_manager(model_nameger: ModelManager):
model, model_path = model_nameger.fetch_model("translator", require_model_path=True)
def from_model_manager(model_manager: ModelManager):
model, model_path = model_manager.fetch_model("translator", require_model_path=True)
translator = Translator(tokenizer_path=model_path, model=model)
return translator