gokaygokay commited on
Commit
ca1953f
1 Parent(s): 546b6c5

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +3 -1
README.md CHANGED
@@ -12,6 +12,8 @@ pipeline_tag: text2text-generation
12
  ```python
13
  from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
14
 
 
 
15
  # Model checkpoint
16
  model_checkpoint = "gokaygokay/Flux-Prompt-Enhance"
17
 
@@ -21,7 +23,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
21
  # Model
22
  model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)
23
 
24
- enhancer = pipeline('text2text-generation', model=model, tokenizer=tokenizer, repetition_penalty=1.5, device="cuda")
25
  max_target_length = 256
26
  prefix = "enhance prompt: "
27
 
 
12
  ```python
13
  from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
14
 
15
+ device = "cuda" if torch.cuda.is_available() else "cpu"
16
+
17
  # Model checkpoint
18
  model_checkpoint = "gokaygokay/Flux-Prompt-Enhance"
19
 
 
23
  # Model
24
  model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)
25
 
26
+ enhancer = pipeline('text2text-generation', model=model, tokenizer=tokenizer, device=device)
27
  max_target_length = 256
28
  prefix = "enhance prompt: "
29