Update app.py

app.py

@@ -17,7 +17,7 @@ vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
 )
 vl_gpt = vl_gpt.cuda().eval()
 
-# === Image generation function ===
+# === Image generation ===
 def text_to_image_generate(input_prompt, output_path, vl_chat_processor, vl_gpt, temperature=1.0, parallel_size=1, cfg_weight=3.0):
     torch.cuda.empty_cache()
 
@@ -52,13 +52,22 @@ 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, :])
 
@@ -87,7 +96,10 @@ def text_to_image_generate(input_prompt, output_path, vl_chat_processor, vl_gpt,
         Image.fromarray(dec[i]).save(save_path)
         output_images.append(save_path)
 
-    return output_images[:1]  # just return first image
+    torch.cuda.empty_cache()
+    torch.cuda.ipc_collect()
+
+    return output_images[:1]  # return only the first image
 
 
 # === Gradio handler ===
@@ -105,9 +117,9 @@ demo = gr.ChatInterface(
     title="Janus Text-to-Image",
     description="Generate images from natural language prompts using Janus-4o-7B",
     examples=[
-        "a cat wearing a spacesuit on Mars",
-        "a beautiful sunset over the mountains",
-        "a photorealistic dog riding a bicycle"
+        "a cat sitting on a windowsill",
+        "a futuristic city at sunset",
+        "a dragon flying over mountains",
     ],
     theme="soft",
 )