Update app.py

app.py

@@ -52,29 +52,31 @@ flux_generator = FluxGenerator()
 @spaces.GPU
 @torch.inference_mode()
 def generate_image(
-        prompt,
+        prompt_name,
         id_image,
         start_step,
         guidance,
         seed,
-        style,
+        style_name,
         true_cfg,
         width=896,
         height=1152,
         num_steps=20,
         id_weight=1.0,
-        neg_prompt="bad quality, worst quality, text, signature, watermark, extra limbs , nudity ,Blurred face, low-quality details, exaggerated facial expressions, unrealistic skin texture, distorted proportions",
+        neg_prompt="bad quality, worst quality, text, signature, watermark, extra limbs",
         timestep_to_start_cfg=1,
         max_sequence_length=128,
 ):
-    flux_generator.t5.max_length = max_sequence_length
-    prompt , neg_prompt= apply_prompt(prompt, prompt, neg_prompt)
+    # Fetch and apply prompt and style
+    pos_prompt, neg_prompt = apply_prompt(prompt_name)
+    pos_prompt, neg_prompt = apply_style(style_name, pos_prompt, neg_prompt)
+
     seed = int(seed)
     if seed == -1:
         seed = None
 
     opts = SamplingOptions(
-        prompt=prompt,
+        prompt=pos_prompt,
         width=width,
         height=height,
         num_steps=num_steps,
@@ -84,6 +86,8 @@ def generate_image(
 
     if opts.seed is None:
         opts.seed = torch.Generator(device="cpu").seed()
+
+    # Log generation details
     print(f"Generating '{opts.prompt}' with seed {opts.seed}")
     t0 = time.perf_counter()
 
@@ -165,12 +169,18 @@ def generate_image(
                 flux_generator.pulid_model.debug_img_list)
 
 def apply_style(style_name: str, positive: str, negative: str = "") -> tuple[str, str]:
-    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
-    return p.replace("{prompt}", positive), n + ' ' + negative
-
-def apply_prompt(prompt: str, positive: str, negative: str = "") -> tuple[str, str]:
-    p,n = prompt_dict.get(prompt,prompt_dict[DEFAULT_PROMPT_NAME])
-    return p.replace("{prompt}", positive), n + ' ' + negative
+    # Fetch the style from the styles dict
+    style_data = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    pos_style = style_data['prompt'].replace("{prompt}", positive)
+    neg_style = style_data['negative_prompt'] + ' ' + negative
+    return pos_style, neg_style
+
+def apply_prompt(prompt_name: str, positive: str = "", negative: str = "") -> tuple[str, str]:
+    # Fetch the prompt from the prompt_dict
+    prompt_data = prompt_dict.get(prompt_name, prompt_dict[DEFAULT_PROMPT_NAME])
+    pos_prompt = prompt_data['prompt'].replace("{prompt}", positive)
+    neg_prompt = prompt_data['negative_prompt'] + ' ' + negative
+    return pos_prompt, neg_prompt