commit

app.py

@@ -413,7 +413,8 @@ CHECKPOINTS = ESRGANUpscalerCheckpoints(
 
 device = DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 DTYPE = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
-enhancer = ESRGANUpscaler(checkpoints=CHECKPOINTS, device="cpu", dtype=DTYPE)
+enhancer = ESRGANUpscaler(checkpoints=CHECKPOINTS, device=DEVICE, dtype=DTYPE)
+enhancer.to(DEVICE)
 
 # logging
 
@@ -434,8 +435,8 @@ pegasus_name = "google/pegasus-xsum"
 # precision data
 
 seq=512
-width=1500
-height=1500
+width=1536
+height=1536
 image_steps=8
 img_accu=0
 
@@ -554,7 +555,7 @@ def upscaler(
     return enhanced_image
 
 def summarize_text(
-    text, max_length=30, num_beams=4, early_stopping=True
+    text, max_length=10, num_beams=4, early_stopping=True
 ):
     log(f'CALL summarize_text')
     summary = pegasus_tokenizer.decode( pegasus_model.generate(
@@ -570,8 +571,7 @@ def generate_random_string(length):
     characters = str(ascii_letters + digits)
     return ''.join(random.choice(characters) for _ in range(length))
 
-@spaces.GPU(duration=180)
-def pipe_generate(p1,p2):
+def pipe_generate_image(p1,p2):
     log(f'CALL pipe_generate')
     imgs = image_pipe(
             prompt=p1,
@@ -579,7 +579,7 @@ def pipe_generate(p1,p2):
             height=height,
             width=width,
             guidance_scale=img_accu,
-            num_images_per_prompt=6,
+            num_images_per_prompt=1,
             num_inference_steps=image_steps,
             max_sequence_length=seq,
             generator=torch.Generator(device).manual_seed(int(str(random.random()).split(".")[1]))
@@ -612,6 +612,20 @@ def add_song_cover_text(img,artist,song):
 
     return img
 
+@spaces.GPU(duration=180)
+def all_pipes(img,pos,neg,artist,song):
+    imgs = pipe_generate_image(pos,neg)
+
+    names = []
+    index = 1
+    for img in imgs:
+        labeled_img = add_song_cover_text(img,artist,song)
+        enhanced_img = upscaler(labeled_img)
+        name = f'{pos_artist} - {pos_song} ({index}).png'
+        enhanced_img.save(name)
+        names.append(name)
+    return names
+
 def handle_generate(artist,song,genre,lyrics):
 
     log(f'CALL handle_generate')
@@ -631,19 +645,9 @@ def handle_generate(artist,song,genre,lyrics):
 
         Negative: {neg}
     """)
-    
-    imgs = pipe_generate(pos,neg)
-
-    names = []
-    index = 1
-    for img in imgs:
-        labeled_img = add_song_cover_text(img,pos_artist,pos_song)
-        enhanced_img = upscaler(labeled_img)
-        name = f'{pos_artist} - {pos_song} ({index}).png'
-        enhanced_img.save(name)
-        names.append(name)
-    return names
 
+    return all_pipes(img,pos,neg,pos_artist,pos_song)
+    
 # entry
 
 if __name__ == "__main__":
@@ -682,17 +686,12 @@ if __name__ == "__main__":
             run = gr.Button("Generate",elem_classes="btn")
 
             with gr.Row():
-                cover1 = gr.Image(interactive=False,container=False,elem_classes="image-container", label="Result", show_label=True, type='filepath', show_share_button=False)
-                cover2 = gr.Image(interactive=False,container=False,elem_classes="image-container", label="Result", show_label=True, type='filepath', show_share_button=False)
-                cover3 = gr.Image(interactive=False,container=False,elem_classes="image-container", label="Result", show_label=True, type='filepath', show_share_button=False)
-                cover4 = gr.Image(interactive=False,container=False,elem_classes="image-container", label="Result", show_label=True, type='filepath', show_share_button=False)
-                cover5 = gr.Image(interactive=False,container=False,elem_classes="image-container", label="Result", show_label=True, type='filepath', show_share_button=False)
-                cover6 = gr.Image(interactive=False,container=False,elem_classes="image-container", label="Result", show_label=True, type='filepath', show_share_button=False)
+                cover = gr.Image(interactive=False,container=False,elem_classes="image-container", label="Result", show_label=True, type='filepath', show_share_button=False)
 
         run.click(
             fn=handle_generate,
             inputs=[artist,song,genre,lyrics],
-            outputs=[cover1,cover2,cover3,cover4,cover5,cover6]
+            outputs=[cover]
         )
 
     demo.queue().launch()