gokaygokay commited on
Commit
01f7732
·
verified ·
1 Parent(s): 9e88c26

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -25,8 +25,8 @@ model_path = snapshot_download(
25
  )
26
 
27
  # VLM Captioner
28
- vlm_model = PaliGemmaForConditionalGeneration.from_pretrained("gokaygokay/sd3-long-captioner").to(device).eval()
29
- vlm_processor = PaliGemmaProcessor.from_pretrained("gokaygokay/sd3-long-captioner")
30
 
31
  # Prompt Enhancer
32
  enhancer_medium = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device=device)
@@ -45,7 +45,7 @@ def create_captions_rich(image):
45
  input_len = model_inputs["input_ids"].shape[-1]
46
 
47
  with torch.inference_mode():
48
- generation = vlm_model.generate(**model_inputs, max_new_tokens=256, do_sample=False)
49
  generation = generation[0][input_len:]
50
  decoded = vlm_processor.decode(generation, skip_special_tokens=True)
51
 
 
25
  )
26
 
27
  # VLM Captioner
28
+ vlm_model = PaliGemmaForConditionalGeneration.from_pretrained("gokaygokay/sd3-long-captioner-v2").to(device).eval()
29
+ vlm_processor = PaliGemmaProcessor.from_pretrained("gokaygokay/sd3-long-captioner-v2")
30
 
31
  # Prompt Enhancer
32
  enhancer_medium = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device=device)
 
45
  input_len = model_inputs["input_ids"].shape[-1]
46
 
47
  with torch.inference_mode():
48
+ generation = vlm_model.generate(**model_inputs, repetition_penalty=1.10, max_new_tokens=256, do_sample=False)
49
  generation = generation[0][input_len:]
50
  decoded = vlm_processor.decode(generation, skip_special_tokens=True)
51