gokaygokay

/

Flux-Prompt-Enhance

like 39

import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

device = "cuda" if torch.cuda.is_available() else "cpu"

# Model checkpoint
model_checkpoint = "gokaygokay/Flux-Prompt-Enhance"

# Tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

# Model
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint).to(device)

max_target_length = 256
prefix = "enhance prompt: "

def enhance_prompt(prompt, max_length=max_target_length):
    # Set a random seed for this generation
    seed = torch.randint(0, 2**32 - 1, (1,)).item()
    torch.manual_seed(seed)
    
    input_text = prefix + prompt
    input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)
    
    outputs = model.generate(
        input_ids,
        max_length=max_length,
        num_return_sequences=1,
        do_sample=True,
        temperature=0.7,
        repetition_penalty=1.2
    )
    
    return seed, tokenizer.decode(outputs[0], skip_special_tokens=True)

short_prompt = "beautiful house with text 'hello'"
seed, final_answer = enhance_prompt(short_prompt)
print(f"Seed: {seed}")
print(final_answer)

You can add do_sample=True and make the temperature higher. Kolay gelsin. 👍