Update app.py

app.py

@@ -10,8 +10,8 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 device = torch.device("cuda:0")
 
-llm = AutoModelForCausalLM.from_pretrained("Azure99/blossom-v5-14b", torch_dtype=torch.float16, device_map="auto")
-tokenizer = AutoTokenizer.from_pretrained("Azure99/blossom-v5-14b")
+llm = AutoModelForCausalLM.from_pretrained("Azure99/blossom-v5-4b", torch_dtype=torch.float16, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained("Azure99/blossom-v5-4b")
 diffusion_pipe = DiffusionPipeline.from_pretrained(
     "playgroundai/playground-v2.5-1024px-aesthetic",
     torch_dtype=torch.float16,
@@ -34,7 +34,7 @@ def save_image(img):
 
 
 LLM_PROMPT = '''你的任务是从输入的[作画要求]中抽取画面描述(description)，然后description翻译为英文(en_description)，最后对en_description进行扩写(expanded_description)，增加足够多的细节，且符合人类的第一直觉。
-[输出]是一个json，包含description、en_description、expanded_description三个字符串字段，请直接输出json，不要输出任何 无关内容。
+[输出]是一个json，包含description、en_description、expanded_description三个字符串字段，请直接输出一个完整的json，不要输出任何解释或其他无关内容。
 
 下面是一些示例：
 [作画要求]->"画一幅画：落霞与孤鹜齐飞，秋水共长天一色。"
@@ -60,22 +60,29 @@ def generate(
         prompt: str,
         progress=gr.Progress(track_tqdm=True),
 ):
-    
     input_ids = get_input_ids(LLM_PROMPT.replace("$USER_PROMPT", json.dumps(prompt, ensure_ascii=False)), BOT_PREFIX)
     generation_kwargs = dict(input_ids=torch.tensor([input_ids]).to(llm.device), do_sample=True,
                              max_new_tokens=512, temperature=0.5, top_p=0.85, top_k=50, repetition_penalty=1.05)
     llm_result = llm.generate(**generation_kwargs)
     llm_result = llm_result.cpu()[0][len(input_ids):]
     llm_result = BOT_PREFIX + tokenizer.decode(llm_result, skip_special_tokens=True)
+    print("----------")
+    print(prompt)
     print(llm_result)
-    expanded_description = json.loads(llm_result)["expanded_description"]
-    print(expanded_description)
+    en_prompt = prompt
+    expanded_prompt = prompt
+    try:
+        en_prompt = json.loads(llm_result)["en_description"]
+        expanded_prompt = json.loads(llm_result)["expanded_description"]
+    except:
+        print("error, fallback to original prompt")
+        pass
 
     seed = random.randint(0, 2147483647)
     generator = torch.Generator().manual_seed(seed)
 
     images = diffusion_pipe(
-        prompt=expanded_description,
+        prompt=[en_prompt, expanded_prompt],
         negative_prompt=None,
         width=1024,
         height=1024,
@@ -107,7 +114,7 @@ with gr.Blocks(css=css) as demo:
                 container=False,
             )
             run_button = gr.Button("Run", scale=0)
-        result = gr.Gallery(label="Result", columns=1, show_label=False)
+        result = gr.Gallery(label="Result", columns=2, rows=1, show_label=False)
 
     gr.on(
         triggers=[