app.py

import gradio as gr
import torch
from transformers import AutoConfig, AutoModelForCausalLM
from janus.models import MultiModalityCausalLM, VLChatProcessor
from janus.utils.io import load_pil_images
from PIL import Image
import numpy as np
import os
import time
import spaces

# 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)
if torch.cuda.is_available():
    vl_gpt = vl_gpt.to(torch.bfloat16).cuda()
else:
    vl_gpt = vl_gpt.to(torch.float16)

vl_chat_processor = VLChatProcessor.from_pretrained(model_path)
tokenizer = vl_chat_processor.tokenizer
cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'

@torch.inference_mode()
@spaces.GPU(duration=120) 
def multimodal_understanding(image, question, seed, top_p, temperature):
    # Clear CUDA cache before generating
    torch.cuda.empty_cache()
    
    # set seed
    torch.manual_seed(seed)
    np.random.seed(seed)
    torch.cuda.manual_seed(seed)
    
    conversation = [
        {
            "role": "<|User|>",
            "content": f"<image_placeholder>\n{question}",
            "images": [image],
        },
        {"role": "<|Assistant|>", "content": ""},
    ]
    
    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)
    
    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=4000,
        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

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):
    # 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):
    # 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 = 5
    
    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)]

# Gradio interface with improved UI
with gr.Blocks(theme=gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="indigo",
)) as demo:
    gr.Markdown(
        """
        # Deepseek Multimodal
        ### Advanced AI for Visual Understanding and Generation
        
        This powerful multimodal AI system combines:
        * **Visual Analysis**: Advanced image understanding and medical image interpretation
        * **Creative Generation**: High-quality image generation from text descriptions
        * **Interactive Chat**: Natural conversation about visual content
        """
    )
    
    with gr.Tabs():
        # Visual Chat Tab
        with gr.Tab("Visual Understanding", icon="image"):
            with gr.Row(equal_height=True):
                with gr.Column(scale=1):
                    image_input = gr.Image(
                        label="Upload Image",
                        type="numpy",
                        elem_classes="image-preview"
                    )
                    
                with gr.Column(scale=1):
                    question_input = gr.Textbox(
                        label="Question or Analysis Request",
                        placeholder="Ask a question about the image or request detailed analysis...",
                        lines=3
                    )
                    with gr.Row():
                        und_seed_input = gr.Number(
                            label="Seed",
                            precision=0,
                            value=42,
                            container=False
                        )
                        top_p = gr.Slider(
                            minimum=0,
                            maximum=1,
                            value=0.95,
                            step=0.05,
                            label="Top-p",
                            container=False
                        )
                        temperature = gr.Slider(
                            minimum=0,
                            maximum=1,
                            value=0.1,
                            step=0.05,
                            label="Temperature",
                            container=False
                        )
                    
                    understanding_button = gr.Button(
                        "Analyze Image",
                        variant="primary",
                        size="lg"
                    )
            
            understanding_output = gr.Textbox(
                label="Analysis Results",
                lines=10,
                show_copy_button=True
            )
            
            gr.Examples(
                label="Medical Analysis Examples",
                examples=[
                    [
                        """You are an AI assistant trained to analyze medical images. Analyze the attached fundus photograph in extreme detail, following a structured approach analogous to an ophthalmologist's examination. Provide a differential diagnosis solely based on this image analysis, without assuming any clinical history. Structure your response as follows:

Analysis Methodology (Concise): List, very briefly, the key anatomical areas/features you will assess, in the order of assessment (e.g., Optic Disc, Vessels, Macula, Periphery, Overall Quality). Do not describe the analysis process in detail here – just name the areas.
Detailed Image Analysis & Percentage Breakdown: Analyze the image, addressing each area listed in Step 1. For each area:
Provide a highly detailed, objective description, using precise ophthalmological terminology. Quantify observations whenever possible (e.g., cup-to-disc ratio, A/V ratio).
State the percentage of that area you were able to confidently analyze, based on image quality and clarity. For example: "Optic Disc: 90% analyzable (10% obscured by slight blurring at the superior margin)." "Macula: 100% analyzable." "Peripheral Retina (Nasal): 60% analyzable (40% not visible in the image)." Be precise.
For any areas where analysis is incomplete (<100%), briefly explain the limiting factor (e.g., poor focus, limited field of view, artifact).
Differential Diagnosis (Image-Based Only): Based solely on your Step 2 analysis, provide:
Most Likely Diagnosis (from image findings).
Other Possible Diagnoses (from image findings).
Rationale: For each diagnosis, briefly link specific image findings to the diagnostic criteria.""",
                        "fundus.webp",
                    ],
                ],
                inputs=[question_input, image_input],
            )

        # Image Generation Tab
        with gr.Tab("Image Generation", icon="wand"):
            with gr.Column():
                prompt_input = gr.Textbox(
                    label="Image Description",
                    placeholder="Describe the image you want to create in detail...",
                    lines=3
                )
                
                with gr.Row():
                    cfg_weight_input = gr.Slider(
                        minimum=1,
                        maximum=10,
                        value=5,
                        step=0.5,
                        label="Guidance Scale",
                        info="Higher values create images that more closely match your prompt"
                    )
                    t2i_temperature = gr.Slider(
                        minimum=0,
                        maximum=1,
                        value=1.0,
                        step=0.05,
                        label="Temperature",
                        info="Controls randomness in generation"
                    )
                    seed_input = gr.Number(
                        label="Seed (Optional)",
                        precision=0,
                        value=12345,
                        info="Set for reproducible results"
                    )
                
                generation_button = gr.Button(
                    "Generate Images",
                    variant="primary",
                    size="lg"
                )
                
                image_output = gr.Gallery(
                    label="Generated Images",
                    columns=3,
                    rows=2,
                    height=500,
                    object_fit="contain"
                )
                
                gr.Examples(
                    label="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 glass of red wine on a reflective surface.",
                        "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."
                    ],
                    inputs=prompt_input,
                )

    # Connect components
    understanding_button.click(
        multimodal_understanding,
        inputs=[image_input, question_input, und_seed_input, top_p, temperature],
        outputs=understanding_output
    )
    
    generation_button.click(
        fn=generate_image,
        inputs=[prompt_input, seed_input, cfg_weight_input, t2i_temperature],
        outputs=image_output
    )

# Add custom CSS
demo.load(css="""
    .gradio-container {
        font-family: 'Inter', -apple-system, sans-serif;
    }
    .image-preview {
        min-height: 300px;
        max-height: 500px;
        width: 100%;
        object-fit: contain;
        border-radius: 8px;
        border: 2px solid #eee;
    }
    .tabs.tab-nav {
        border-bottom: 2px solid #eee;
        margin-bottom: 2rem;
    }
    .tab-nav {
        background: white;
        padding: 1rem;
        border-radius: 8px;
        box-shadow: 0 2px 4px rgba(0,0,0,0.05);
    }
    .examples-table {
        font-size: 0.9rem;
    }
    .gr-button.gr-button-lg {
        padding: 12px 24px;
        font-size: 1.1rem;
    }
    .gr-input, .gr-select {
        border-radius: 6px;
    }
    .gr-form {
        background: white;
        padding: 20px;
        border-radius: 12px;
        box-shadow: 0 4px 6px rgba(0,0,0,0.05);
    }
    .gr-panel {
        border: none;
        background: transparent;
    }
    .footer {
        text-align: center;
        margin-top: 2rem;
        padding: 1rem;
        color: #666;
    }
""")

# Launch the demo
if __name__ == "__main__":
    demo.launch(share=True)