import gradio as gr
import torch
from PIL import Image
import numpy as np
from transformers import BlipProcessor, BlipForConditionalGeneration

device = "cuda" if torch.cuda.is_available() else "cpu"

def load_models():
    print(f"Loading Enhanced BLIP Large on {device}...")
    processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
    model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large").to(device)
    return processor, model

processor, model = load_models()

def describe_logic(image):
    if image is None: return "Please upload an image."
    image_pil = Image.fromarray(image).convert('RGB') if isinstance(image, np.ndarray) else image.convert('RGB')
    
    inputs = processor(image_pil, return_tensors="pt").to(device)
    
    # High-quality generation parameters
    out = model.generate(
        **inputs, 
        max_new_tokens=100, 
        num_beams=5, 
        repetition_penalty=1.2, 
        length_penalty=1.0
    )
    return processor.decode(out[0], skip_special_tokens=True)

def chat_logic(message, history):
    return "Chat connected. Using Studio9 Architecture."

with gr.Blocks(theme='glass') as demo:
    gr.Markdown("# 🌌 Studio9: Enhanced AI Studio")
    with gr.Tabs():
        with gr.TabItem("💬 Chat"):
            gr.ChatInterface(fn=chat_logic, type='messages')
        with gr.TabItem("✨ Vision"):
            with gr.Row():
                with gr.Column():
                    img_in = gr.Image(type='numpy')
                    btn = gr.Button("Generate Detailed Caption", variant='primary')
                with gr.Column():
                    out = gr.Textbox(label="Detailed Description", lines=8)
            btn.click(describe_logic, img_in, out)

if __name__ == '__main__':
    demo.launch()