Update app.py

app.py

@@ -11,12 +11,15 @@ import time
 from Upsample import RealESRGAN
 import spaces  # Import spaces for ZeroGPU compatibility
 
+
 # Load model and processor
 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:
@@ -30,11 +33,10 @@ 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)
 
-last_uploaded_image = None
-
 @torch.inference_mode()
 @spaces.GPU(duration=120) 
-def multimodal_understanding(image, question, seed, top_p, temperature, progress=gr.Progress(track_tqdm=True)):
+# Multimodal Chat function with conversation history
+def multimodal_chat(image, message, chat_history, seed, top_p, temperature, progress=gr.Progress(track_tqdm=True)):
     # Clear CUDA cache before generating
     torch.cuda.empty_cache()
     
@@ -43,56 +45,295 @@ def multimodal_understanding(image, question, seed, top_p, temperature, progress
     np.random.seed(seed)
     torch.cuda.manual_seed(seed)
     
-    conversation = [
-        {
-            "role": "<|User |>",
-            "content": f"<image_placeholder>\n{question}",
-            "images": [image],
-        },
-        {"role": "<|Assistant|>", "content": ""},
-    ]
+    # Process the conversation history and add current message
+    conversation = []
+    
+    # Check if we have existing history
+    if chat_history:
+        # Add previous conversation turns
+        for user_msg, assistant_msg in chat_history:
+            conversation.append({
+                "role": "<|User|>",
+                "content": user_msg,
+                "images": [],  # No images for previous turns
+            })
+            conversation.append({
+                "role": "<|Assistant|>", 
+                "content": assistant_msg,
+            })
+    
+    # Add the current user message with image (if provided)
+    user_content = message
+    images_list = []
+    
+    # Only include image placeholder if image is provided or this is the first message
+    if image is not None:
+        user_content = f"<image_placeholder>\n{message}"
+        images_list = [image]
+        
+    conversation.append({
+        "role": "<|User|>",
+        "content": user_content,
+        "images": images_list,
+    })
+    conversation.append({"role": "<|Assistant|>", "content": ""})
+    
+    # Process images (if any)
+    pil_images = []
+    if image is not None:
+        pil_images = [Image.fromarray(image)]
     
-    pil_images = [Image.fromarray(image)]
-    prepare_inputs = vl_chat_processor(conversations=conversation, images=pil_images, force_batchify=True
-                                      ).to(cuda_device, dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16)
+    prepare_inputs = vl_chat_processor(
+        conversations=conversation, images=pil_images, force_batchify=True
+    ).to(cuda_device, dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16)
     
     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, temperature=temperature, top_p=top_p,
-                                             do_sample=False if temperature == 0 else True, use_cache=True,)
+    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,
+    )
     
     answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
-    return answer
+    
+    # Update chat history
+    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)):
+    # Clear CUDA cache before generating
+    torch.cuda.empty_cache()
+    
+    tokens = torch.zeros((parallel_size * 2, len(input_ids)), dtype=torch.int).to(cuda_device)
+    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).to(cuda_device)
+
+    pkv = None
+    for i in range(image_token_num_per_image):
+        with torch.no_grad():
+            outputs = vl_gpt.language_model.model(inputs_embeds=inputs_embeds,
+                                                use_cache=True,
+                                                past_key_values=pkv)
+            pkv = outputs.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)
+            next_token = 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)
+            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])
+
+    return generated_tokens.to(dtype=torch.int), patches
+
+def unpack(dec, width, height, parallel_size=5):
+    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, width, height, 3), dtype=np.uint8)
+    visual_img[:, :, :] = dec
+
+    return visual_img
+
+
+
+@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)):
+    # Clear CUDA cache and avoid tracking gradients
+    torch.cuda.empty_cache()
+    # Set the seed for reproducible results
+    if seed is not None:
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        np.random.seed(seed)
+    width = 384
+    height = 384
+    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='')
+        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)]
+        stime = time.time()
+        ret_images = [image_upsample(Image.fromarray(images[i])) for i in range(parallel_size)]
+        print(f'upsample time: {time.time() - stime}')
+        return ret_images
+
+
+@spaces.GPU(duration=60)
+def image_upsample(img: Image.Image) -> Image.Image:
+    if img is None:
+        raise Exception("Image not uploaded")
+    
+    width, height = img.size
+    
+    if width >= 5000 or height >= 5000:
+        raise Exception("The image is too large.")
+
+    global sr_model
+    result = sr_model.predict(img.convert('RGB'))
+    return result
+
+
+# Helper function to add uploaded image to the chat context
+def add_image_to_chat(image, chat_history):
+    return image, chat_history
+
+
+# Helper function to clear chat history but maintain the image
+def clear_chat(image):
+    return [], image
+
 
 # Gradio interface
 css = '''
 .gradio-container {max-width: 960px !important}
 '''
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("# Janus Pro 7B Chat Interface")
-    
-    chat_history = gr.Chatbot(label="Chat History")
-    message_input = gr.Textbox(label="Type your message here...")
-    image_input = gr.Image(label="Upload an image (optional)")
-    
-    def respond(message, image):
-        global last_uploaded_image
-        if image is not None:
-            last_uploaded_image = image  # Update the last uploaded image
-            response = multimodal_understanding(image, message, seed=42, top_p=0.95, temperature=0.1)
-        elif last_uploaded_image is not None:
-            response = multimodal_understanding(last_uploaded_image, message, seed=42, top_p=0.95, temperature=0.1)
-        else:
-            response = "Please provide an image for multimodal understanding."
+    gr.Markdown("# Janus Pro 7B")
+    
+    with gr.Tab("Multimodal Chat"):
+        gr.Markdown(value="## Multimodal Chat")
         
-        return message, response
+        # 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")
 
-    def submit_message(message, image):
-        response = respond(message, image)
-        return message, response
+        prompt_input = gr.Textbox(label="Prompt. (Prompt in more detail can help produce better images!)")
+    
+        generation_button = gr.Button("Generate Images")
+    
+        image_output = gr.Gallery(label="Generated Images", columns=4, rows=1)
 
-    message_input.submit(submit_message, inputs=[message_input, image_input], outputs=[message_input, chat_history])
+        with gr.Accordion("Advanced options", open=False):
+            with gr.Row():
+                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,
+        )
+    
+    # Chat interface interactions
+    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]
+    )
+    
+    clear_button.click(
+        clear_chat,
+        inputs=[current_image],
+        outputs=[chat_interface, current_image]
+    )
+    
+    # T2I interface interactions
+    generation_button.click(
+        fn=generate_image,
+        inputs=[prompt_input, seed_input, cfg_weight_input, t2i_temperature],
+        outputs=image_output
+    )
 
 demo.launch(share=True)