Update app.py

app.py

@@ -4,12 +4,10 @@ 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  # Import spaces for ZeroGPU compatibility
-
+import spaces
 
 # Load model and processor
 model_path = "deepseek-ai/Janus-Pro-7B"
@@ -30,7 +28,6 @@ cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 @torch.inference_mode()
 @spaces.GPU(duration=120) 
-# Multimodal Understanding function
 def multimodal_understanding(image, question, seed, top_p, temperature):
     # Clear CUDA cache before generating
     torch.cuda.empty_cache()
@@ -54,7 +51,6 @@ def multimodal_understanding(image, question, seed, top_p, temperature):
         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(
@@ -73,7 +69,6 @@ def multimodal_understanding(image, question, seed, top_p, temperature):
     answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
     return answer
 
-
 def generate(input_ids,
              width,
              height,
@@ -113,8 +108,6 @@ def generate(input_ids,
             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])
 
@@ -129,8 +122,6 @@ def unpack(dec, width, height, parallel_size=5):
 
     return visual_img
 
-
-
 @torch.inference_mode()
 @spaces.GPU(duration=120)  # Specify a duration to avoid timeout
 def generate_image(prompt,
@@ -169,28 +160,82 @@ def generate_image(prompt,
                         parallel_size=parallel_size)
 
         return [Image.fromarray(images[i]).resize((768, 768), Image.LANCZOS) for i in range(parallel_size)]
-        
 
-# Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown(value="# AutoMultimodal")
-    with gr.Row():
-        image_input = gr.Image()
-        with gr.Column():
-            question_input = gr.Textbox(label="Question")
-            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")
+# 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
         
-    understanding_button = gr.Button("Chat")
-    understanding_output = gr.Textbox(label="Response")
-
-    examples_inpainting = gr.Examples(
-        label="Multimodal Medical Prompts",
-        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:
+        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:
@@ -200,43 +245,74 @@ For any areas where analysis is incomplete (<100%), briefly explain the limiting
 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.
+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
+                )
                 
-"""      ,          "fundus.webp",
-            ],
-        ],
-        inputs=[question_input, image_input],
-    )
-    
-        
-    gr.Markdown(value="# Text-to-Image Multimodal Generation")
+                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,
+                )
 
-    
-    
-    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")
-
-    prompt_input = gr.Textbox(label="Prompt. (Prompt in more detail can help produce better images!)")
-    seed_input = gr.Number(label="Seed (Optional)", precision=0, value=12345)
-
-    generation_button = gr.Button("Generate Images")
-
-    image_output = gr.Gallery(label="Generated Images", columns=2, rows=2, height=300)
-
-    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 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.\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,
-    )
-    
+    # Connect components
     understanding_button.click(
         multimodal_understanding,
         inputs=[image_input, question_input, und_seed_input, top_p, temperature],
@@ -249,5 +325,58 @@ Rationale: For each diagnosis, briefly link specific image findings to the diagn
         outputs=image_output
     )
 
-demo.launch(share=True)
-# demo.queue(concurrency_count=1, max_size=10).launch(server_name="0.0.0.0", server_port=37906, root_path="/path")
+# 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)
+