commit

app.py

@@ -31,7 +31,7 @@ from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file, save_file
 from diffusers import FluxPipeline
 from PIL import Image, ImageDraw, ImageFont
-from transformers import PegasusForConditionalGeneration, PegasusTokenizerFast
+from transformers import pipeline
 from refiners.fluxion.utils import manual_seed
 from refiners.foundationals.latent_diffusion import Solver, solvers
 from refiners.foundationals.latent_diffusion.stable_diffusion_1.multi_upscaler import (
@@ -502,9 +502,6 @@ image_pipe.enable_model_cpu_offload()
 image_pipe.enable_vae_slicing()
 image_pipe.enable_vae_tiling()
 
-pegasus_tokenizer = PegasusTokenizerFast.from_pretrained("google/pegasus-xsum")
-pegasus_model = PegasusForConditionalGeneration.from_pretrained("google/pegasus-xsum")
-
 # functionality
 
 def upscaler(
@@ -554,10 +551,8 @@ def summarize_text(
     text, max_length=16
 ):
     log(f'CALL summarize_text')
-    summary = pegasus_tokenizer.decode( pegasus_model.generate(
-        pegasus_tokenizer(text,return_tensors="pt").input_ids,
-        max_length=max_length
-    )[0], skip_special_tokens=True)
+    summ = pipeline("summarization", model="google/pegasus-xsum")
+    summary = summ("summarize: " + text)[0]["summary_text"]
     log(f'RET summarize_text with summary as {summary}')
     return summary