Update app.py

app.py

@@ -18,7 +18,7 @@ vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
 vl_gpt = vl_gpt.cuda().eval()
 
 # === 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):
+def text_to_image_generate(input_prompt, output_path, vl_chat_processor, vl_gpt, temperature=1.0, parallel_size=2, cfg_weight=5.0):
     torch.cuda.empty_cache()
 
     conversation = [
@@ -26,72 +26,82 @@ def text_to_image_generate(input_prompt, output_path, vl_chat_processor, vl_gpt,
         {"role": "<|Assistant|>", "content": ""},
     ]
 
-    prompt = vl_chat_processor.apply_sft_template_for_multi_turn_prompts(
+    sft_format = vl_chat_processor.apply_sft_template_for_multi_turn_prompts(
         conversations=conversation,
         sft_format=vl_chat_processor.sft_format,
         system_prompt="",
-    ) + vl_chat_processor.image_start_tag
+    )
+    prompt = sft_format + vl_chat_processor.image_start_tag
 
-    image_token_num = 576
+    image_token_num_per_image = 576
     img_size = 384
     patch_size = 16
 
-    input_ids = tokenizer.encode(prompt)
-    input_ids = torch.LongTensor(input_ids)
-    tokens = torch.zeros((parallel_size * 2, len(input_ids)), dtype=torch.int).cuda()
-
-    for i in range(parallel_size * 2):
-        tokens[i, :] = input_ids
-        if i % 2 != 0:
-            tokens[i, 1:-1] = tokenizer.pad_token_id
-
-    inputs_embeds = vl_gpt.language_model.get_input_embeddings()(tokens)
-    generated_tokens = torch.zeros((parallel_size, image_token_num), dtype=torch.int).cuda()
-
-    for i in range(image_token_num):
-        if i == 0:
-            current_inputs = inputs_embeds
-        else:
-            current_inputs = img_embeds.unsqueeze(1)
-
-        # ✅ No past_key_values, crash-safe
-        outputs = vl_gpt.language_model.model(
-            inputs_embeds=current_inputs,
-            use_cache=False
+    with torch.inference_mode():
+        input_ids = vl_chat_processor.tokenizer.encode(prompt)
+        input_ids = torch.LongTensor(input_ids)
+        tokens = torch.zeros((parallel_size * 2, len(input_ids)), dtype=torch.int).cuda()
+
+        for i in range(parallel_size * 2):
+            tokens[i, :] = input_ids
+            if i % 2 != 0:
+                tokens[i, 1:-1] = vl_chat_processor.pad_id
+
+        inputs_embeds = vl_gpt.language_model.get_input_embeddings()(tokens)
+        generated_tokens = torch.zeros((parallel_size, image_token_num_per_image), dtype=torch.int).cuda()
+        
+        past_key_values = None
+        
+        for i in range(image_token_num_per_image):
+            outputs = vl_gpt.language_model.model(
+                inputs_embeds=inputs_embeds, 
+                use_cache=True, 
+                past_key_values=past_key_values
+            )
+            
+            hidden_states = outputs.last_hidden_state
+            logits = vl_gpt.gen_head(hidden_states[:, -1, :])
+            
+            logit_cond = logits[0::2, :]
+            logit_uncond = logits[1::2, :]
+            
+            logits = logit_uncond + cfg_weight * (logit_cond - logit_uncond)
+            probs = torch.softmax(logits / temperature, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+            generated_tokens[:, i] = next_token.squeeze(dim=-1)
+            
+            # Prepare next token for both conditional and unconditional
+            next_token_expanded = torch.cat([next_token.unsqueeze(dim=1), next_token.unsqueeze(dim=1)], dim=1).view(-1)
+            img_embeds = vl_gpt.prepare_gen_img_embeds(next_token_expanded)
+            inputs_embeds = img_embeds.unsqueeze(dim=1)
+            
+            # Update past_key_values for next iteration
+            past_key_values = outputs.past_key_values
+
+        # Decode generated tokens to images
+        dec = vl_gpt.gen_vision_model.decode_code(
+            generated_tokens.to(dtype=torch.int), 
+            shape=[parallel_size, 8, img_size // patch_size, img_size // patch_size]
         )
-
-        hidden_states = outputs.last_hidden_state
-        logits = vl_gpt.gen_head(hidden_states[:, -1, :])
-
-        logit_cond = logits[0::2]
-        logit_uncond = logits[1::2]
-        logits = logit_uncond + cfg_weight * (logit_cond - logit_uncond)
-        probs = torch.softmax(logits / temperature, dim=-1)
-
-        next_token = torch.multinomial(probs, num_samples=1)
-        generated_tokens[:, i] = next_token.squeeze(dim=-1)
-
-        next_token = torch.cat([next_token, next_token], dim=1).reshape(-1)
-        img_embeds = vl_gpt.prepare_gen_img_embeds(next_token)
-
-    dec = vl_gpt.gen_vision_model.decode_code(
-        generated_tokens.to(dtype=torch.int),
-        shape=[parallel_size, 8, img_size // patch_size, img_size // patch_size]
-    )
-    dec = dec.to(torch.float32).detach().cpu().numpy().transpose(0, 2, 3, 1)
-    dec = np.clip((dec + 1) / 2 * 255, 0, 255).astype(np.uint8)
-
-    os.makedirs(os.path.dirname(output_path), exist_ok=True)
-    output_images = []
-    for i in range(parallel_size):
-        save_path = output_path.replace('.png', f'_{i}.png')
-        Image.fromarray(dec[i]).save(save_path)
-        output_images.append(save_path)
+        dec = dec.to(torch.float32).cpu().numpy().transpose(0, 2, 3, 1)
+        dec = np.clip((dec + 1) / 2 * 255, 0, 255)
+        
+        visual_img = np.zeros((parallel_size, img_size, img_size, 3), dtype=np.uint8)
+        visual_img[:, :, :] = dec
+
+        # Create output directory
+        output_dir = os.path.dirname(output_path)
+        if output_dir:
+            os.makedirs(output_dir, exist_ok=True)
+        
+        output_images = []
+        for i in range(parallel_size):
+            save_path = output_path.replace('.png', f'_{i}.png')
+            Image.fromarray(visual_img[i]).save(save_path)
+            output_images.append(save_path)
 
     torch.cuda.empty_cache()
-    torch.cuda.ipc_collect()
-
-    return output_images[:1]  # return only the first image
+    return output_images
 
 
 # === Gradio handler ===
@@ -99,6 +109,7 @@ def text_to_image_generate(input_prompt, output_path, vl_chat_processor, vl_gpt,
 def janus_generate_image(message, history):
     output_path = "./output/image.png"
     images = text_to_image_generate(message, output_path, vl_chat_processor, vl_gpt)
+    # Return the first generated image
     return {"role": "assistant", "content": {"path": images[0]}}