Update app.py

app.py

@@ -52,21 +52,13 @@ def text_to_image_generate(input_prompt, output_path, vl_chat_processor, vl_gpt,
     for i in range(image_token_num):
         if i == 0:
             outputs = vl_gpt.language_model.model(inputs_embeds=inputs_embeds, use_cache=True)
-            past_key_values = outputs.past_key_values
         else:
-            # 🧠 SAFE DETACH TRICK
-            with torch.no_grad():
-                past_key_values = tuple(
-                    tuple(pkv.detach() for pkv in layer)
-                    for layer in past_key_values
-                )
-
             outputs = vl_gpt.language_model.model(
                 inputs_embeds=img_embeds.unsqueeze(1),
                 use_cache=True,
                 past_key_values=past_key_values
             )
-            past_key_values = outputs.past_key_values
+        past_key_values = outputs.past_key_values
 
         hidden_states = outputs.last_hidden_state
         logits = vl_gpt.gen_head(hidden_states[:, -1, :])