Update app.py

app.py

@@ -87,6 +87,22 @@ class calculateDuration:
         else:
             print(f"Elapsed time: {elapsed:.6f} seconds")
 
+##############################
+# ===== Helper: truncate_prompt =====
+##############################
+def truncate_prompt(prompt: str) -> str:
+    """
+    Uses the global pipeline's tokenizer (assumed available as `pipe.tokenizer`)
+    to truncate the prompt to the maximum allowed length.
+    """
+    try:
+        tokenizer = pipe.tokenizer
+        tokenized = tokenizer(prompt, truncation=True, max_length=tokenizer.model_max_length, return_tensors="pt")
+        return tokenizer.decode(tokenized.input_ids[0], skip_special_tokens=True)
+    except Exception as e:
+        print(f"Error in truncate_prompt: {e}")
+        return prompt
+
 ##############################
 # ===== enhance.py =====
 ##############################
@@ -137,7 +153,6 @@ def generate(message, max_new_tokens=256, temperature=0.9, top_p=0.95, repetitio
 ##############################
 # ===== lora_handling.py =====
 ##############################
-# Default LoRA list for initial UI setup
 loras = [
     {"image": "placeholder.jpg", "title": "Placeholder LoRA", "repo": "placeholder/repo", "weights": None, "trigger_word": ""}
 ]
@@ -163,7 +178,6 @@ def flux_pipe_call_that_returns_an_iterable_of_images(self,
                                                        good_vae: Optional[Any] = None):
     height = height or self.default_sample_size * self.vae_scale_factor
     width = width or self.default_sample_size * self.vae_scale_factor
-    
     self.check_inputs(
         prompt,
         prompt_2,
@@ -341,6 +355,8 @@ def prepare_prompt(prompt: str, selected_index: Optional[int], loras_list: list)
             prompt_mash = f"{prompt} {trigger_word}"
     else:
         prompt_mash = prompt
+    # Truncate the prompt using the tokenizer to ensure token indices are in range.
+    prompt_mash = truncate_prompt(prompt_mash)
     return prompt_mash
 
 def unload_lora_weights(pipe, pipe_i2i):
@@ -405,10 +421,10 @@ class ModelManager:
             tokenizer_2=self.pipe.tokenizer_2,
             torch_dtype=DTYPE,
         ).to(DEVICE)
-        # Bind custom LoRA method to the pipeline class (to avoid __slots__ issues)
+        # Bind the custom LoRA method to the pipeline class (to avoid __slots__ issues)
         self.pipe.__class__.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images
-    
-    @spaces.GPU(duration=100) 
+
+    @spaces.GPU(duration=100)
     def generate_image(self, prompt_mash, steps, seed, cfg_scale, width, height, lora_scale):
         generator = torch.Generator(device=DEVICE).manual_seed(seed)
         with calculateDuration("Generating image"):
@@ -425,6 +441,7 @@ class ModelManager:
             ):
                 yield img
 
+    @spaces.GPU(duration=100)
     def generate_image_to_image(self, prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, lora_scale, seed):
         generator = torch.Generator(device=DEVICE).manual_seed(seed)
         image_input = load_image_from_path(image_input_path)