Update app.py

app.py

@@ -3,9 +3,89 @@ import gradio as gr
 import numpy as np
 import spaces
 import torch
-from diffusers import AutoPipelineForText2Image, AutoencoderKL #,EulerDiscreteScheduler
+from diffusers import AutoPipelineForText2Image, AutoencoderKL
 from compel import Compel, ReturnedEmbeddingsType
 
+import re
+
+def tokenize_line(text, tokenizer):
+    tokens = tokenizer.tokenize(text)
+    return tokens
+
+def parse_prompt_attention(text):
+    res = []
+    pattern = re.compile(r"\(([^)]+):([\d\.]+)\)")
+    matches = pattern.findall(text)
+    for match in matches:
+        res.append((match[0], float(match[1])))
+    return res
+
+def prompt_attention_to_invoke_prompt(attention_list):
+    prompt = ""
+    for item in attention_list:
+        prompt += f"({item[0]}:{item[1]}) "
+    return prompt.strip()
+
+def merge_embeds(prompts, compel):
+    embeds = []
+    pooled_embeds = []
+    for prompt in prompts:
+        conditioning, pooled = compel(prompt)
+        embeds.append(conditioning)
+        pooled_embeds.append(pooled)
+    # 合并嵌入，这里使用平均值，可以根据需要调整
+    merged_embed = torch.mean(torch.stack(embeds), dim=0)
+    merged_pooled = torch.mean(torch.stack(pooled_embeds), dim=0)
+    return merged_embed, merged_pooled
+
+def get_embed_new(prompt, pipeline, compel, only_convert_string=False, compel_process_sd=False):
+    if compel_process_sd:
+        return merge_embeds(tokenize_line(prompt, pipeline.tokenizer), compel)
+    else:
+        # fix bug weights conversion excessive emphasis
+        prompt = prompt.replace("((", "(").replace("))", ")")
+
+    # Convert to Compel
+    attention = parse_prompt_attention(prompt)
+    global_attention_chunks = []
+
+    for att in attention:
+        for chunk in att[0].split(','):
+            temp_prompt_chunks = tokenize_line(chunk, pipeline.tokenizer)
+            for small_chunk in temp_prompt_chunks:
+                temp_dict = {
+                    "weight": round(att[1], 2),
+                    "length": len(pipeline.tokenizer.tokenize(f'{small_chunk},')),
+                    "prompt": f'{small_chunk},'
+                }
+                global_attention_chunks.append(temp_dict)
+
+    max_tokens = pipeline.tokenizer.model_max_length - 2
+    global_prompt_chunks = []
+    current_list = []
+    current_length = 0
+    for item in global_attention_chunks:
+        if current_length + item['length'] > max_tokens:
+            global_prompt_chunks.append(current_list)
+            current_list = [[item['prompt'], item['weight']]]
+            current_length = item['length']
+        else:
+            if not current_list:
+                current_list.append([item['prompt'], item['weight']])
+            else:
+                if item['weight'] != current_list[-1][1]:
+                    current_list.append([item['prompt'], item['weight']])
+                else:
+                    current_list[-1][0] += f" {item['prompt']}"
+            current_length += item['length']
+    if current_list:
+        global_prompt_chunks.append(current_list)
+
+    if only_convert_string:
+        return ' '.join([prompt_attention_to_invoke_prompt(i) for i in global_prompt_chunks])
+
+    return merge_embeds([prompt_attention_to_invoke_prompt(i) for i in global_prompt_chunks], compel)
+
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>你现在运行在CPU上 但是此项目只支持GPU.</p>"
 
@@ -21,7 +101,6 @@ if torch.cuda.is_available():
         use_safetensors=True,
         add_watermarker=False
     )
-    #pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
     pipe.to("cuda")
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
@@ -45,14 +124,28 @@ def infer(
 ):
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator().manual_seed(seed)
-    compel = Compel(tokenizer=[pipe.tokenizer, pipe.tokenizer_2] , text_encoder=[pipe.text_encoder, pipe.text_encoder_2], returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, requires_pooled=[False, True])
-    conditioning, pooled = compel(prompt)
-    
+    # 初始化 Compel 实例
+    compel_instance = Compel(
+        tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
+        text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
+        returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+        requires_pooled=[False, True]
+    )
+    # 使用 get_embed_new 获取嵌入
+    conditioning, pooled = get_embed_new(prompt, pipe, compel_instance)
+
+    # 处理反向提示（negative_prompt）
+    if use_negative_prompt and negative_prompt:
+        negative_conditioning, negative_pooled = get_embed_new(negative_prompt, pipe, compel_instance)
+    else:
+        negative_conditioning = None
+        negative_pooled = None
+
     image = pipe(
-        #prompt=prompt,
         prompt_embeds=conditioning,
         pooled_prompt_embeds=pooled,
-        negative_prompt=negative_prompt,
+        negative_prompt_embeds=negative_conditioning,
+        negative_pooled_prompt_embeds=negative_pooled,
         width=width,
         height=height,
         guidance_scale=guidance_scale,
@@ -60,6 +153,7 @@ def infer(
         generator=generator,
         use_resolution_binning=use_resolution_binning,
     ).images[0]
+    image.save("output_image.png")
     return image, seed
 
 examples = [
@@ -74,7 +168,7 @@ footer {
     visibility: hidden
 }
 '''
-    
+
 with gr.Blocks(css=css) as demo:
     gr.Markdown("""# 梦羽的模型生成器
         ### 快速生成NoobAIXL v0.5的模型图片 V1.0模型在另一个项目上""")
@@ -145,7 +239,7 @@ with gr.Blocks(css=css) as demo:
         outputs=[result, seed],
         fn=infer
     )
-    
+
     use_negative_prompt.change(
         fn=lambda x: gr.update(visible=x),
         inputs=use_negative_prompt,
@@ -153,7 +247,7 @@ with gr.Blocks(css=css) as demo:
     )
 
     gr.on(
-        triggers=[prompt.submit,run_button.click],
+        triggers=[prompt.submit, run_button.click],
         fn=infer,
         inputs=[
             prompt,