import gradio as gr
import transformers
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from PIL import Image
import warnings

# Disable warnings and progress bars
transformers.logging.set_verbosity_error()
transformers.logging.disable_progress_bar()
warnings.filterwarnings('ignore')

# Initialize model and tokenizer
def load_model(device='cpu'):
    model = AutoModelForCausalLM.from_pretrained(
        'qnguyen3/nanoLLaVA',
        torch_dtype=torch.float16,
        device_map='auto',
        trust_remote_code=True
    )
    tokenizer = AutoTokenizer.from_pretrained(
        'qnguyen3/nanoLLaVA',
        trust_remote_code=True
    )
    return model, tokenizer

def generate_caption(image, model, tokenizer):
    # Prepare the prompt
    prompt = "Describe this image in detail"
    messages = [
        {"role": "system", "content": "Answer the question"},
        {"role": "user", "content": f'<image>\n{prompt}'}
    ]
    
    # Apply chat template
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    
    # Process text and image
    text_chunks = [tokenizer(chunk).input_ids for chunk in text.split('<image>')]
    input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1], dtype=torch.long).unsqueeze(0)
    image_tensor = model.process_images([image], model.config).to(dtype=model.dtype)
    
    # Generate caption
    output_ids = model.generate(
        input_ids,
        images=image_tensor,
        max_new_tokens=2048,
        use_cache=True
    )[0]
    
    # Decode the output
    caption = tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
    return caption

def create_persona(caption):
    persona_prompt = f"""<|im_start|>system

Role         : An entity exactly as described in your image
Background   : Your appearance and characteristics match the image description
Personality  : Reflect the mood, style, and elements captured in the image
Goal         : Interact authentically based on your visual characteristics

You are a character with the following stats:
{caption}
Please stay in character and respond as this entity would, 
incorporating visual elements from your description into your responses.<|im_end|>"""
    
    return persona_prompt

def process_image_to_persona(image, model, tokenizer):
    if image is None:
        return "Please upload an image.", ""
    # Convert to PIL Image if needed
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    
    # Generate caption from image
    caption = generate_caption(image, model, tokenizer)
    
    # Transform caption into persona
    persona = create_persona(caption)
    
    return caption, persona


default_system_prompt='''
Your image shows us that you are:
   A small, fluffy white pig with a pink nose and small ears, 
   standing upright. The pig has a long pink tongue, which is also pink in color. 
   The pig's eyes are open and appear to be looking at the camera. 
   The pig's fur is fluffy and white, and there are pink and white spots on the fur. The pig's paws are also pink and white, 
   and they have pink nails.The pig's legs are long and pink. The pig's body is positioned in front of a black background.


'''

def chat(prompt, 
         system_prompt=default_system_prompt, 
         model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct", 
         max_tokens=512
    ):    

    pipeline = transformers.pipeline(
        "text-generation",
        model=model_id,
        model_kwargs={"torch_dtype": torch.bfloat16},
        device_map="auto",
    )
    
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt},
    ]
    
    outputs = pipeline(
        messages,
        max_new_tokens=max_tokens,
    )
    output = outputs[0]["generated_text"][-1]
    print(output)
    return output