Update app.py

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 import torch
 from langs import LANGS
-import os
 
 import numpy as np
 import modin.pandas as pd
@@ -14,7 +13,7 @@ device1 = "cuda" if torch.cuda.is_available() else "cpu"
 pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", safety_checker=None)
 pipe = pipe.to(device1)
 
-pipe1 = pipeline('text-generation', model='RamAnanth1/distilgpt2-sd-prompts')
+pipe1 = pipeline('text-generation', model='daspartho/prompt-extend')
 
 TASK = "translation"
 CKPT = "facebook/nllb-200-distilled-600M"
@@ -46,23 +45,8 @@ def translate(text):
 
 #prompt
 
-stable_diffusion = gr.Blocks.load(name="spaces/runwayml/stable-diffusion-v1-5")
-clip_interrogator_2 = gr.Blocks.load(name="spaces/fffiloni/CLIP-Interrogator-2")
-
-def get_images(prompt):
-    gallery_dir = stable_diffusion(prompt, fn_index=2)    
-    img_results = [os.path.join(gallery_dir, img) for img in os.listdir(gallery_dir)]  
-    return img_results[0]
-
-def get_new_prompt(img, mode):
-    interrogate = clip_interrogator_2(img, mode, 12, api_name="clipi2")
-    return interrogate
-
-def infer(input):
-    prompt = pipe1(input+',', num_return_sequences=1)[0]["generated_text"]
-    img = get_images(prompt) 
-    result = get_new_prompt(img, 'fast')    
-    return result[0]
+def extend_prompt(prompt):
+    return pipe1(prompt+',', num_return_sequences=1)[0]["generated_text"]
 
 #stable diffusion
 
@@ -93,7 +77,7 @@ with gr.Blocks() as demo:
         btn1 = gr.Button("让您的提示词更详细一些吧")
         out2 = gr.Textbox(label = "更详细的英文提示词(可编辑此内容)", lines=1)
     
-    btn1.click(infer, [inp2], [out2])
+    btn1.click(extend_prompt, [inp2], [out2])
 
     with gr.Row().style(equal_height=True):
         inp3 = out2