Update app.py

app.py

@@ -14,40 +14,45 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 
 @spaces.GPU(duration=120)
-def generate_response(image, prompt, max_length, temperature):
+def generate_description(image, detail_level):
+    if image is None:
+        return "Please upload an image to generate a description."
+    
+    image = Image.open(image).convert("RGB")
+    
+    detail_prompts = {
+        "Brief": "Provide a brief description of this image in 2-3 sentences.",
+        "Detailed": "Describe this image in detail, including main subjects, colors, composition, and any notable elements.",
+        "Comprehensive": "Provide a comprehensive analysis of this image, including subjects, colors, composition, mood, potential symbolism, and any other relevant details you can observe."
+    }
+    
+    prompt = detail_prompts[detail_level]
+    
     messages = [
-        {"role": "system", "content": "You are a helpful assistant that can analyze images and text."},
+        {"role": "system", "content": "You are a highly observant AI assistant specialized in describing images accurately and in detail."},
         {"role": "user", "content": prompt}
     ]
     formatted_prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
 
-    # Preprocess the image
-    if image is not None:
-        image = Image.open(image).convert("RGB")
-        inputs = tokenizer(formatted_prompt, images=[image], return_tensors="pt", padding=True).to(model.device)
-    else:
-        inputs = tokenizer(formatted_prompt, return_tensors="pt", padding=True).to(model.device)
+    inputs = tokenizer(formatted_prompt, images=[image], return_tensors="pt", padding=True).to(model.device)
     
-    # Generate
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=max_length,
+            max_new_tokens=300,
             do_sample=True,
-            temperature=temperature,
-            top_k=100,
+            temperature=0.7,
+            top_k=50,
             top_p=0.95,
         )
     
-    # Decode and return the response
-    response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
-    return response
+    description = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+    return description.strip()
 
 # Custom CSS
 css = """
 body {
-    background-color: #1a1a2e;
-    color: #e0e0e0;
+    background-color: #f0f0f5;
     font-family: 'Arial', sans-serif;
 }
 .container {
@@ -56,37 +61,34 @@ body {
     padding: 20px;
 }
 .gradio-container {
-    background-color: #16213e;
+    background-color: white;
     border-radius: 15px;
     box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
 }
 .header {
-    background-color: #0f3460;
+    background-color: #4a90e2;
+    color: white;
     padding: 20px;
     border-radius: 15px 15px 0 0;
     text-align: center;
     margin-bottom: 20px;
 }
 .header h1 {
-    color: #e94560;
     font-size: 2.5em;
     margin-bottom: 10px;
 }
-.header p {
-    color: #a0a0a0;
-}
 .input-group, .output-group {
-    background-color: #1a1a2e;
+    background-color: #f9f9f9;
     padding: 20px;
     border-radius: 10px;
     margin-bottom: 20px;
 }
 .input-group label, .output-group label {
-    color: #e94560;
+    color: #4a90e2;
     font-weight: bold;
 }
 .generate-btn {
-    background-color: #e94560 !important;
+    background-color: #4a90e2 !important;
     color: white !important;
     border: none !important;
     border-radius: 5px !important;
@@ -96,72 +98,31 @@ body {
     transition: background-color 0.3s ease !important;
 }
 .generate-btn:hover {
-    background-color: #c81e45 !important;
-}
-.example-prompts {
-    background-color: #1f2b47;
-    padding: 15px;
-    border-radius: 10px;
-    margin-bottom: 20px;
-}
-.example-prompts h3 {
-    color: #e94560;
-    margin-bottom: 10px;
-}
-.example-prompts ul {
-    list-style-type: none;
-    padding-left: 0;
-}
-.example-prompts li {
-    margin-bottom: 5px;
-    cursor: pointer;
-    transition: color 0.3s ease;
-}
-.example-prompts li:hover {
-    color: #e94560;
+    background-color: #3a7bc8 !important;
 }
 """
 
-# Example prompts
-example_prompts = [
-    "Describe this image in detail.",
-    "What emotions does this image evoke?",
-    "Imagine a story based on this image.",
-    "What technical aspects of photography are demonstrated in this image?",
-    "How might this image be used in advertising?"
-]
-
 # Gradio interface
 with gr.Blocks(css=css) as iface:
     gr.HTML(
         """
         <div class="header">
-            <h1>Pixtral-12B Multimodal Generation</h1>
-            <p>Generate text responses based on images and prompts using the powerful Pixtral-12B model.</p>
+            <h1>Pixtral Image Description Generator</h1>
+            <p>Upload an image and get a detailed description using the powerful Pixtral-12B model.</p>
         </div>
         """
     )
 
     with gr.Group():
-        with gr.Group(elem_classes="example-prompts"):
-            gr.HTML("<h3>Example Prompts:</h3>")
-            example_buttons = [gr.Button(prompt) for prompt in example_prompts]
-
         with gr.Group(elem_classes="input-group"):
-            image_input = gr.Image(type="filepath", label="Upload an image (optional)")
-            prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...", lines=5)
-            max_length = gr.Slider(minimum=1, maximum=500, value=128, step=1, label="Max Length")
-            temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
-            generate_btn = gr.Button("Generate", elem_classes="generate-btn")
+            image_input = gr.Image(type="filepath", label="Upload an image")
+            detail_level = gr.Radio(["Brief", "Detailed", "Comprehensive"], label="Description Detail Level", value="Detailed")
+            generate_btn = gr.Button("Generate Description", elem_classes="generate-btn")
 
         with gr.Group(elem_classes="output-group"):
-            output = gr.Textbox(label="Generated Text", lines=10)
-
-    generate_btn.click(generate_response, inputs=[image_input, prompt, max_length, temperature], outputs=output)
+            output = gr.Textbox(label="Generated Description", lines=10)
 
-    # Set up example prompt buttons
-    for button in example_buttons:
-        button.click(lambda x: x, inputs=[button], outputs=[prompt])
+    generate_btn.click(generate_description, inputs=[image_input, detail_level], outputs=output)
 
 # Launch the app
 iface.launch()