Update generator.py

generator.py

@@ -261,13 +261,23 @@ class RetroArtConverter:
             return Image.fromarray(depth_colored)
 
     
-    def add_trigger_word(self, prompt):
+    # --- START FIX: Updated function to use lora_choice ---
+    def add_trigger_word(self, prompt, lora_choice="RetroArt"):
         """Add trigger word to prompt if not present"""
-        if TRIGGER_WORD.lower() not in prompt.lower():
+        
+        # Get the correct trigger word from the config dictionary
+        trigger = TRIGGER_WORD.get(lora_choice, TRIGGER_WORD["RetroArt"])
+        
+        if not trigger:
+            return prompt
+            
+        if trigger.lower() not in prompt.lower():
             if not prompt or not prompt.strip():
-                return TRIGGER_WORD
-            return f"{TRIGGER_WORD}, {prompt}"
+                return trigger
+            # Prepend the trigger word as requested
+            return f"{trigger}, {prompt}"
         return prompt
+    # --- END FIX ---
     
     def extract_multi_scale_face(self, face_crop, face):
         """
@@ -548,8 +558,9 @@ class RetroArtConverter:
                     expression_control_scale
                 )
         
-        # Add trigger word
-        prompt = self.add_trigger_word(prompt)
+        # --- START FIX: Pass lora_choice to add_trigger_word ---
+        prompt = self.add_trigger_word(prompt, lora_choice)
+        # --- END FIX ---
         
         # Calculate optimal size with flexible aspect ratio support
         original_width, original_height = input_image.size
@@ -662,7 +673,7 @@ class RetroArtConverter:
                         facial_attrs = get_facial_attributes(face)
                         
                         # Update prompt with detected attributes
-                        prompt = build_enhanced_prompt(prompt, facial_attrs, TRIGGER_WORD)
+                        prompt = build_enhanced_prompt(prompt, facial_attrs, TRIGGER_WORD[lora_choice])
                         
                         # Legacy output for compatibility
                         age = facial_attrs['age']