import gradio as gr
import torch
import os
from PIL import Image
import numpy as np
from transformers import BlipProcessor, BlipForConditionalGeneration, pipeline
from huggingface_hub import InferenceClient

# Handle device selection for local vision models
device = "cuda" if torch.cuda.is_available() else "cpu"

# 1. Vision Model (Local)
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)

# 2. Sentiment Analysis (Local Pipeline)
sentiment_analyzer = pipeline("sentiment-analysis", device=0 if torch.cuda.is_available() else -1)

# 3. Inference API Client
hf_token = os.getenv("HF_TOKEN")
client = InferenceClient(token=hf_token)

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)
    out = model.generate(**inputs, max_new_tokens=50)
    return processor.decode(out[0], skip_special_tokens=True)

def analyze_text(text):
    if not text: return "Enter text to analyze."
    result = sentiment_analyzer(text)[0]
    return f"Label: {result['label']} | Score: {result['score']:.4f}"

def chat_logic(message, history, model_name):
    if not hf_token: return "Error: HF_TOKEN not found in Secrets."
    try:
        messages = [{"role": "user", "content": message}]
        response = ""
        for message in client.chat_completion(model=model_name, messages=messages, max_tokens=500, stream=True):
            token = message.choices[0].delta.content
            if token: response += token
        return response
    except Exception as e: return f"Inference Error: {str(e)}"

def generate_image(prompt):
    if not hf_token: return None
    try:
        return client.text_to_image(prompt, model="stabilityai/stable-diffusion-xl-base-1.0")
    except Exception: return None

with gr.Blocks(theme='glass') as demo:
    gr.Markdown("# 🌌 AI Ultimate Studio v3.6")
    with gr.Tabs():
        with gr.TabItem("💬 Chat"):
            model_choice = gr.Dropdown(choices=["deepseek-ai/DeepSeek-R1-Distill-Llama-8B", "google/gemma-2-9b-it"], value="deepseek-ai/DeepSeek-R1-Distill-Llama-8B", label="Model")
            gr.ChatInterface(fn=chat_logic, additional_inputs=[model_choice], type="messages")
        with gr.TabItem("🎨 Image Gen"):
            with gr.Row():
                with gr.Column():
                    prompt_in = gr.Textbox(label="Prompt")
                    gen_btn = gr.Button("Generate")
                with gr.Column():
                    img_out = gr.Image(label="Result")
            gen_btn.click(generate_image, prompt_in, img_out)
        with gr.TabItem("✨ Vision"):
            img_input = gr.Image(type="numpy")
            describe_btn = gr.Button("Describe")
            text_output = gr.Textbox(label="Result")
            describe_btn.click(describe_logic, img_input, text_output)
        with gr.TabItem("📊 Text Analysis"):
            txt_input = gr.Textbox(label="Sentiment Analysis")
            analyze_btn = gr.Button("Analyze")
            sentiment_output = gr.Textbox(label="Result")
            analyze_btn.click(analyze_text, txt_input, sentiment_output)

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