Update app.py

app.py

@@ -8,6 +8,7 @@ from PIL import Image
 import numpy as np
 import os
 import time
+import re
 from Upsample import RealESRGAN
 import spaces  # Import spaces for ZeroGPU compatibility
 
@@ -17,9 +18,7 @@ model_path = "deepseek-ai/Janus-Pro-7B"
 config = AutoConfig.from_pretrained(model_path)
 language_config = config.language_config
 language_config._attn_implementation = 'eager'
-vl_gpt = AutoModelForCausalLM.from_pretrained(model_path,
-                                             language_config=language_config,
-                                             trust_remote_code=True)
+vl_gpt = AutoModelForCausalLM.from_pretrained(model_path, language_config=language_config, trust_remote_code=True)
 if torch.cuda.is_available():
     vl_gpt = vl_gpt.to(torch.bfloat16).cuda()
 else:
@@ -33,13 +32,80 @@ cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'
 sr_model = RealESRGAN(torch.device('cuda' if torch.cuda.is_available() else 'cpu'), scale=2)
 sr_model.load_weights(f'weights/RealESRGAN_x2.pth', download=False)
 
+# Patterns for detecting image generation requests
+GENERATION_PATTERNS = [r"generate (.+)", r"create (.+)", r"draw (.+)", r"make (.+)", r"show (.+)", r"visualize (.+)", r"imagine (.+)", r"picture (.+)",]
+
+def is_generation_request(message):
+    """Determine if a message is requesting image generation"""
+    message = message.lower().strip()
+    
+    # Check if message explicitly mentions image generation
+    for pattern in GENERATION_PATTERNS:
+        match = re.match(pattern, message, re.IGNORECASE)
+        if match:
+            return True, match.group(1)
+    
+    # Check for specific keywords suggesting image generation
+    image_keywords = ["image", "picture", "photo", "artwork", "illustration", "painting", "drawing"]
+    generation_verbs = ["generate", "create", "make", "produce", "show me", "draw"]
+    
+    for verb in generation_verbs:
+        for keyword in image_keywords:
+            if f"{verb} {keyword}" in message or f"{verb} an {keyword}" in message or f"{verb} a {keyword}" in message:
+                # Extract the prompt (everything after the keyword)
+                pattern = f"{verb}\\s+(?:an?\\s+)?{keyword}\\s+(?:of|showing|depicting|with)?\\s*(.*)"
+                match = re.search(pattern, message, re.IGNORECASE)
+                if match and match.group(1):
+                    return True, match.group(1)
+                else:
+                    # If we can't extract a specific prompt, use the whole message
+                    return True, message
+    
+    return False, None
+
+
 @torch.inference_mode()
 @spaces.GPU(duration=120) 
-# Multimodal Chat function with conversation history
-def multimodal_chat(image, message, chat_history, seed, top_p, temperature, progress=gr.Progress(track_tqdm=True)):
+# Unified chat function that handles both image understanding and generation
+def unified_chat(image, message, chat_history, seed, top_p, temperature, cfg_weight, t2i_temperature, progress=gr.Progress(track_tqdm=True)):
     # Clear CUDA cache before generating
     torch.cuda.empty_cache()
     
+    # Check if this is an image generation request
+    is_gen_request, extracted_prompt = is_generation_request(message)
+    
+    if is_gen_request:
+        # Prepare a more detailed prompt by considering context from the conversation
+        context_prompt = extracted_prompt
+        
+        # Optionally, enhance the prompt with context from previous messages
+        if chat_history and len(chat_history) > 0:
+            # Get the last few turns of conversation for context (limit to last 3 turns)
+            recent_context = chat_history[-3:] if len(chat_history) > 3 else chat_history
+            context_text = " ".join([f"{user}: {user_msg}" for user_msg, _ in recent_context])
+            
+            # Only use context if it's not too long
+            if len(context_text) < 200:  # Arbitrary length limit
+                context_prompt = f"{context_text}. {extracted_prompt}"
+        
+        # Generate images
+        generated_images = generate_image(
+            prompt=context_prompt,
+            seed=seed,
+            guidance=cfg_weight,
+            t2i_temperature=t2i_temperature
+        )
+        
+        # Create a response that includes the generated images
+        response = f"I've generated the following images based on: '{extracted_prompt}'"
+        
+        # Add the images to the chat as the bot's response
+        chat_history.append((message, response))
+        
+        # Return the message, updated history, maintained image context, and generated images
+        return "", chat_history, image, generated_images
+    
+    # Regular chat flow (no image generation)
     # set seed
     torch.manual_seed(seed)
     np.random.seed(seed)
@@ -89,18 +155,10 @@ def multimodal_chat(image, message, chat_history, seed, top_p, temperature, prog
     
     inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
     
-    outputs = vl_gpt.language_model.generate(
-        inputs_embeds=inputs_embeds,
-        attention_mask=prepare_inputs.attention_mask,
-        pad_token_id=tokenizer.eos_token_id,
-        bos_token_id=tokenizer.bos_token_id,
-        eos_token_id=tokenizer.eos_token_id,
-        max_new_tokens=512,
-        do_sample=False if temperature == 0 else True,
-        use_cache=True,
-        temperature=temperature,
-        top_p=top_p,
-    )
+    outputs = vl_gpt.language_model.generate(inputs_embeds=inputs_embeds, attention_mask=prepare_inputs.attention_mask,
+                                             pad_token_id=tokenizer.eos_token_id, bos_token_id=tokenizer.bos_token_id,
+                                             eos_token_id=tokenizer.eos_token_id, max_new_tokens=512, temperature=temperature, top_p=top_p,
+                                             do_sample=False if temperature == 0 else True, use_cache=True,)
     
     answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
     
@@ -108,18 +166,11 @@ def multimodal_chat(image, message, chat_history, seed, top_p, temperature, prog
     chat_history.append((message, answer))
     
     # Keep the last uploaded image in context
-    return "", chat_history, image
-
-
-def generate(input_ids,
-             width,
-             height,
-             temperature: float = 1,
-             parallel_size: int = 5,
-             cfg_weight: float = 5,
-             image_token_num_per_image: int = 576,
-             patch_size: int = 16,
-             progress=gr.Progress(track_tqdm=True)):
+    return "", chat_history, image, None
+
+
+def generate(input_ids, width, height, temperature: float = 1, parallel_size: int = 5, cfg_weight: float = 5,
+             image_token_num_per_image: int = 1024, patch_size: int = 16, progress=gr.Progress(track_tqdm=True)):
     # Clear CUDA cache before generating
     torch.cuda.empty_cache()
     
@@ -152,7 +203,7 @@ def generate(input_ids,
             inputs_embeds = img_embeds.unsqueeze(dim=1)
 
     patches = vl_gpt.gen_vision_model.decode_code(generated_tokens.to(dtype=torch.int),
-                                                 shape=[parallel_size, 8, width // patch_size, height // patch_size])
+                                                  shape=[parallel_size, 8, width // patch_size, height // patch_size])
 
     return generated_tokens.to(dtype=torch.int), patches
 
@@ -168,12 +219,8 @@ def unpack(dec, width, height, parallel_size=5):
 
 
 @torch.inference_mode()
-@spaces.GPU(duration=120)  # Specify a duration to avoid timeout
-def generate_image(prompt,
-                   seed=None,
-                   guidance=5,
-                   t2i_temperature=1.0,
-                   progress=gr.Progress(track_tqdm=True)):
+@spaces.GPU(duration=180)  # Specify a duration to avoid timeout
+def generate_image(prompt, seed=None, guidance=5, t2i_temperature=1.0, progress=gr.Progress(track_tqdm=True)):
     # Clear CUDA cache and avoid tracking gradients
     torch.cuda.empty_cache()
     # Set the seed for reproducible results
@@ -181,31 +228,22 @@ def generate_image(prompt,
         torch.manual_seed(seed)
         torch.cuda.manual_seed(seed)
         np.random.seed(seed)
-    width = 384
-    height = 384
+    width = 512
+    height = 512
     parallel_size = 4
     
     with torch.no_grad():
         messages = [{'role': '<|User|>', 'content': prompt},
                     {'role': '<|Assistant|>', 'content': ''}]
         text = vl_chat_processor.apply_sft_template_for_multi_turn_prompts(conversations=messages,
-                                                                   sft_format=vl_chat_processor.sft_format,
-                                                                   system_prompt='')
+                                                                           sft_format=vl_chat_processor.sft_format, system_prompt='')
         text = text + vl_chat_processor.image_start_tag
         
         input_ids = torch.LongTensor(tokenizer.encode(text))
-        output, patches = generate(input_ids,
-                                   width // 16 * 16,
-                                   height // 16 * 16,
-                                   cfg_weight=guidance,
-                                   parallel_size=parallel_size,
-                                   temperature=t2i_temperature)
-        images = unpack(patches,
-                        width // 16 * 16,
-                        height // 16 * 16,
-                        parallel_size=parallel_size)
-
-        # return [Image.fromarray(images[i]).resize((768, 768), Image.LANCZOS) for i in range(parallel_size)]
+        output, patches = generate(input_ids, width // 16 * 16, height // 16 * 16, cfg_weight=guidance,
+                                   parallel_size=parallel_size, temperature=t2i_temperature)
+        images = unpack(patches, width // 16 * 16, height // 16 * 16, parallel_size=parallel_size)
+
         stime = time.time()
         ret_images = [image_upsample(Image.fromarray(images[i])) for i in range(parallel_size)]
         print(f'upsample time: {time.time() - stime}')
@@ -219,7 +257,7 @@ def image_upsample(img: Image.Image) -> Image.Image:
     
     width, height = img.size
     
-    if width >= 5000 or height >= 5000:
+    if width >= 4096 or height >= 4096:
         raise Exception("The image is too large.")
 
     global sr_model
@@ -234,7 +272,7 @@ def add_image_to_chat(image, chat_history):
 
 # Helper function to clear chat history but maintain the image
 def clear_chat(image):
-    return [], image
+    return [], image, None
 
 
 # Gradio interface
@@ -242,98 +280,80 @@ css = '''
 .gradio-container {max-width: 960px !important}
 '''
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("# Janus Pro 7B")
-    
-    with gr.Tab("Multimodal Chat"):
-        gr.Markdown(value="## Multimodal Chat")
-        
-        # State variables to maintain context
-        chat_history = gr.State([])
-        current_image = gr.State(None)
-        
-        with gr.Row():
-            with gr.Column(scale=1):
-                image_input = gr.Image(label="Upload Image (only needed once)")
-                upload_button = gr.Button("Add Image to Chat")
-                
-                with gr.Accordion("Advanced options", open=False):
-                    und_seed_input = gr.Number(label="Seed", precision=0, value=42)
-                    top_p = gr.Slider(minimum=0, maximum=1, value=0.95, step=0.05, label="top_p")
-                    temperature = gr.Slider(minimum=0, maximum=1, value=0.1, step=0.05, label="temperature")
-                
-                clear_button = gr.Button("Clear Chat")
-            
-            with gr.Column(scale=2):
-                chat_interface = gr.Chatbot(label="Chat History", height=500)
-                message_input = gr.Textbox(label="Your message", placeholder="Ask about the image or continue the conversation...", lines=2)
-                chat_button = gr.Button("Send")
-                
-        examples_chat = gr.Examples(
-            label="Multimodal Chat examples",
-            examples=[
-                [
-                    "explain this meme",
-                    "doge.png",
-                ],
-                [
-                    "Convert the formula into latex code.",
-                    "equation.png",
-                ],
-            ],
-            inputs=[message_input, image_input],
-        )
-    
-    with gr.Tab("Text-to-Image Generation"):
-        gr.Markdown(value="## Text-to-Image Generation")
-
-        prompt_input = gr.Textbox(label="Prompt. (Prompt in more detail can help produce better images!)")
+    gr.Markdown("# Janus Pro 7B - Unified Chat Interface")
     
-        generation_button = gr.Button("Generate Images")
+    # State variables to maintain context
+    chat_history = gr.State([])
+    current_image = gr.State(None)
     
-        image_output = gr.Gallery(label="Generated Images", columns=4, rows=1)
-
-        with gr.Accordion("Advanced options", open=False):
-            with gr.Row():
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_input = gr.Image(label="Upload Image (optional)")
+            upload_button = gr.Button("Add Image to Chat")
+            
+            with gr.Accordion("Chat Options", open=False):
+                und_seed_input = gr.Number(label="Seed", precision=0, value=42)
+                top_p = gr.Slider(minimum=0, maximum=1, value=0.95, step=0.05, label="top_p")
+                temperature = gr.Slider(minimum=0, maximum=1, value=0.1, step=0.05, label="temperature")
+            
+            with gr.Accordion("Image Generation Options", open=False):
                 cfg_weight_input = gr.Slider(minimum=1, maximum=10, value=5, step=0.5, label="CFG Weight")
-                t2i_temperature = gr.Slider(minimum=0, maximum=1, value=1.0, step=0.05, label="temperature")
-            seed_input = gr.Number(label="Seed (Optional)", precision=0, value=1234)
-    
-        examples_t2i = gr.Examples(
-            label="Text to image generation examples.",
-            examples=[
-                "Master shifu racoon wearing drip attire as a street gangster.",
-                "The face of a beautiful girl",
-                "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-                "A cute and adorable baby fox with big brown eyes, autumn leaves in the background enchanting,immortal,fluffy, shiny mane,Petals,fairyism,unreal engine 5 and Octane Render,highly detailed, photorealistic, cinematic, natural colors.",
-                "The image features an intricately designed eye set against a circular backdrop adorned with ornate swirl patterns that evoke both realism and surrealism. At the center of attention is a strikingly vivid blue iris surrounded by delicate veins radiating outward from the pupil to create depth and intensity. The eyelashes are long and dark, casting subtle shadows on the skin around them which appears smooth yet slightly textured as if aged or weathered over time.\n\nAbove the eye, there's a stone-like structure resembling part of classical architecture, adding layers of mystery and timeless elegance to the composition. This architectural element contrasts sharply but harmoniously with the organic curves surrounding it. Below the eye lies another decorative motif reminiscent of baroque artistry, further enhancing the overall sense of eternity encapsulated within each meticulously crafted detail. \n\nOverall, the atmosphere exudes a mysterious aura intertwined seamlessly with elements suggesting timelessness, achieved through the juxtaposition of realistic textures and surreal artistic flourishes. Each component\u2014from the intricate designs framing the eye to the ancient-looking stone piece above\u2014contributes uniquely towards creating a visually captivating tableau imbued with enigmatic allure.",
-            ],
-            inputs=prompt_input,
-        )
+                t2i_temperature_input = gr.Slider(minimum=0, maximum=1, value=1.0, step=0.05, label="Temperature")
+            
+            clear_button = gr.Button("Clear Chat")
+            
+            gr.Markdown("""
+            ### Tips:
+            1. Upload an image to discuss it
+            2. Type commands like "generate [description]" to create images
+            3. Continue chatting about uploaded or generated images
+            4. Use natural language like "show me a sunset" or "create a portrait"
+            """)
+        
+        with gr.Column(scale=2):
+            chat_interface = gr.Chatbot(label="Chat History", height=500)
+            message_input = gr.Textbox(
+                label="Your message", 
+                placeholder="Ask about an image, continue chatting, or generate new images by typing 'generate [description]'",
+                lines=2
+            )
+            chat_button = gr.Button("Send")
+            generated_images = gr.Gallery(label="Generated Images", visible=True, columns=2, rows=2)
     
     # Chat interface interactions
-    upload_button.click(
-        add_image_to_chat,
-        inputs=[image_input, chat_history],
-        outputs=[current_image, chat_history]
-    )
+    upload_button.click(add_image_to_chat, inputs=[image_input, chat_history], outputs=[current_image, chat_history])
     
     chat_button.click(
-        multimodal_chat,
-        inputs=[current_image, message_input, chat_interface, und_seed_input, top_p, temperature],
-        outputs=[message_input, chat_interface, current_image]
+        unified_chat,
+        inputs=[current_image, message_input, chat_interface, und_seed_input, top_p, temperature, cfg_weight_input, t2i_temperature_input],
+        outputs=[message_input, chat_interface, current_image, generated_images]
+    )
+    
+    # Also trigger on Enter key
+    message_input.submit(
+        unified_chat,
+        inputs=[current_image, message_input, chat_interface, und_seed_input, top_p, temperature, cfg_weight_input, t2i_temperature_input],
+        outputs=[message_input, chat_interface, current_image, generated_images]
     )
     
     clear_button.click(
         clear_chat,
         inputs=[current_image],
-        outputs=[chat_interface, current_image]
+        outputs=[chat_interface, current_image, generated_images]
     )
     
-    # T2I interface interactions
-    generation_button.click(
-        fn=generate_image,
-        inputs=[prompt_input, seed_input, cfg_weight_input, t2i_temperature],
-        outputs=image_output
+    # Examples for the unified interface
+    examples = gr.Examples(
+        label="Example queries",
+        examples=[
+            ["What's in this image?"],
+            ["Generate a cute kitten with big eyes"],
+            ["Show me a mountain landscape at sunset"],
+            ["Can you explain what's happening in this picture?"],
+            ["Create an astronaut riding a horse"],
+            ["Generate a futuristic cityscape with flying cars"],
+        ],
+        inputs=message_input,
     )
 
 demo.launch(share=True)