Update app.py

app.py

@@ -1,34 +1,27 @@
 import os
-import PIL.Image
 import torch
 import numpy as np
+from PIL import Image
 import gradio as gr
 from transformers import AutoModelForCausalLM
 from janus.models import MultiModalityCausalLM, VLChatProcessor
 import spaces
 
-# Load model and processor
+# === Load model and processor ===
 model_path = "FreedomIntelligence/Janus-4o-7B"
 vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
 tokenizer = vl_chat_processor.tokenizer
+
 vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
     model_path, trust_remote_code=True, torch_dtype=torch.bfloat16
 )
 vl_gpt = vl_gpt.cuda().eval()
 
-# Define image generation function
-@spaces.GPU(duration=120)
-def janus_generate_image(message, history):
-    prompt = message
-    output_path = "./output_image.png"
-    images = text_to_image_generate(prompt, output_path, vl_chat_processor, vl_gpt, parallel_size=1)
-    return {"role": "assistant", "content": gr.Image(images[0])}
-
-# Optimized text-to-image generation
-
+# === Image generation function ===
 def text_to_image_generate(input_prompt, output_path, vl_chat_processor, vl_gpt, temperature=1.0, parallel_size=1, cfg_weight=5):
     torch.cuda.empty_cache()
 
+    # Apply prompt formatting
     conversation = [
         {"role": "<|User|>", "content": input_prompt},
         {"role": "<|Assistant|>", "content": ""},
@@ -42,28 +35,26 @@ def text_to_image_generate(input_prompt, output_path, vl_chat_processor, vl_gpt,
 
     prompt = sft_format + vl_chat_processor.image_start_tag
     mmgpt = vl_gpt
+
     image_token_num = 576
     img_size = 384
     patch_size = 16
 
     with torch.inference_mode():
-        input_ids = vl_chat_processor.tokenizer.encode(prompt)
+        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] = vl_chat_processor.pad_id
+                tokens[i, 1:-1] = tokenizer.pad_token_id  # More robust
 
         inputs_embeds = mmgpt.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):
-            outputs = mmgpt.language_model.model(
-                inputs_embeds=inputs_embeds,
-                use_cache=False
-            )
+            outputs = mmgpt.language_model.model(inputs_embeds=inputs_embeds, use_cache=False)
             hidden_states = outputs.last_hidden_state
             logits = mmgpt.gen_head(hidden_states[:, -1, :])
             logit_cond = logits[0::2]
@@ -78,6 +69,7 @@ def text_to_image_generate(input_prompt, output_path, vl_chat_processor, vl_gpt,
             img_embeds = mmgpt.prepare_gen_img_embeds(next_token)
             inputs_embeds = img_embeds.unsqueeze(1)
 
+        # Decode image
         dec = mmgpt.gen_vision_model.decode_code(
             generated_tokens.to(dtype=torch.int),
             shape=[parallel_size, 8, img_size // patch_size, img_size // patch_size]
@@ -85,26 +77,33 @@ def text_to_image_generate(input_prompt, output_path, vl_chat_processor, vl_gpt,
         dec = dec.to(torch.float32).cpu().numpy().transpose(0, 2, 3, 1)
         dec = np.clip((dec + 1) / 2 * 255, 0, 255).astype(np.uint8)
 
-        visual_img = np.zeros((parallel_size, img_size, img_size, 3), dtype=np.uint8)
-        visual_img[:, :, :] = dec
-
         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')
-            PIL.Image.fromarray(visual_img[i]).save(save_path)
+            Image.fromarray(dec[i]).save(save_path)
             output_images.append(save_path)
 
         return output_images
 
-# Launch the ChatInterface UI
+# === Gradio handler ===
+@spaces.GPU(duration=120)
+def janus_generate_image(message, history):
+    prompt = message
+    output_path = "./output/image.png"
+    images = text_to_image_generate(prompt, output_path, vl_chat_processor, vl_gpt, parallel_size=1)
+    return {"role": "assistant", "content": images[0]}
+
+# === Gradio UI ===
 demo = gr.ChatInterface(
     fn=janus_generate_image,
     title="Janus Text-to-Image",
     description="Generate images from natural language prompts using Janus-4o-7B",
+    additional_inputs=[],
+    chatbot=gr.Chatbot(show_copy_button=True),
+    examples=["a cat", "a spaceship landing on Mars", "a fantasy castle at sunset"],
     theme="soft",
-    fill_height=True,
-    fill_width=True
 )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()