commit

app.py

@@ -25,6 +25,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
 
 # logging
 
@@ -45,6 +46,7 @@ else:
     device = "cpu"
 
 base = "black-forest-labs/FLUX.1-schnell"
+pegasus_name = "google/pegasus-xsum"
 
 # precision data
 
@@ -118,6 +120,9 @@ function custom(){
 image_pipe = FluxPipeline.from_pretrained(base, torch_dtype=torch.bfloat16).to(device)
 image_pipe.enable_model_cpu_offload()
 
+pegasus_tokenizer = PegasusTokenizerFast.from_pretrained(pegasus_name)
+pegasus_model = PegasusForConditionalGeneration.from_pretrained(pegasus_name)
+
 # functionality
 
 def generate_random_string(length):
@@ -145,8 +150,9 @@ def handle_generate(artist,song,genre,lyrics):
     pos_song = ' '.join(word[0].upper() + word[1:] for word in pos_song.split())
     pos_genre = re.sub(f'[{punctuation}]', '', re.sub("([ \t\n]){1,}", " ", genre)).upper().strip()
     pos_lyrics = re.sub(f'[{punctuation}]', '', re.sub("([ \t\n]){1,}", " ", lyrics)).lower().strip()
+    pos_lyrics_sum = pegasus_tokenizer.decode(pegasus_model.generate(pegasus_tokenizer(pos_lyrics,return_tensors="pt").input_ids)[0], skip_special_tokens=True)
     neg = f"Textual Labeled Distorted Discontinuous Ugly Blurry"
-    pos = f'Realistic Natural Genuine Reasonable Detailed { pos_genre } GENRE SONG COVER FOR { pos_song }: "{ pos_lyrics }"'
+    pos = f'Realistic Natural Genuine Reasonable Detailed, { pos_genre } GENRE SONG - { pos_song }: "{ pos_lyrics_sum }"'
 
     print(f"""
         Positive: {pos}