import gradio as gr
import os, sys
from transformers import AutoModelForCausalLM, LlamaTokenizer
from PIL import Image, UnidentifiedImageError
import torch
import argparse
import time
from pathlib import Path
from itertools import chain

# Initialize the argument parser and parse the command line arguments
parser = argparse.ArgumentParser()
parser.add_argument("--quant", choices=[4, 8, 16], type=int, default=None, help='quantization bits for model loading')
parser.add_argument("--from_pretrained", type=str, default="THUDM/cogvlm-chat-hf", help='pretrained model path')
parser.add_argument("--local_tokenizer", type=str, default="lmsys/vicuna-7b-v1.5", help='local tokenizer path')
parser.add_argument("--fp16", action="store_true", help='use float16 precision')
parser.add_argument("--bf16", action="store_true", help='use bfloat16 precision')
parser.add_argument("--share", action="store_true", help='starts a gradio share')
args = parser.parse_args()

# Setup device and model configurations
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
MODEL_PATH = "THUDM/cogvlm-chat-hf"
tokenizer = LlamaTokenizer.from_pretrained('lmsys/vicuna-7b-v1.5')
torch_type = torch.float16

bit_4 = False
bit_8 = False
bit_32 = False

if(args.quant == 4):
	bit_4=True
	
if(args.quant == 8):
	bit_8=True
	
if(args.quant == 32):
	torch_type = torch.float32	
	

# Load the model based on the arguments
if args.quant in [4, 8]:
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        low_cpu_mem_usage=True,
        load_in_4bit=bit_4,
        load_in_8bit=bit_8,
		torch_dtype=torch_type,
        trust_remote_code=True
    )
if args.quant in [16]:
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        torch_dtype=torch_type,
        low_cpu_mem_usage=True,
        trust_remote_code=True
    ).to(DEVICE)
if args.quant in [32]:
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        torch_dtype=torch_type,
        low_cpu_mem_usage=True,
        trust_remote_code=True
    ).to(DEVICE)

def batch_process(input_folder, input_text, temperature, top_p, top_k, do_sample):
    start_time = time.time()
    processed_count = 0
    error_count = 0
    supported_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff', '.ico', '.webp'}
    
    total_files = sum(1 for f in os.listdir(input_folder) if os.path.splitext(f)[1].lower() in supported_extensions)
    
    for filename in os.listdir(input_folder):
        file_path = os.path.join(input_folder, filename)
        file_extension = os.path.splitext(filename)[1].lower()
        
        if file_extension not in supported_extensions:
            continue
        
        try:
            image = Image.open(file_path).convert('RGB')
            response = process_image(image, input_text, temperature, top_p, top_k, do_sample)
            
            output_path = os.path.splitext(file_path)[0] + '.txt'
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(response)
            
            processed_count += 1
            elapsed_time = time.time() - start_time
            average_time = elapsed_time / processed_count
            print(f"Processed {processed_count}/{total_files}. Remaining: {total_files - processed_count}. Avg. Time: {average_time:.2f} sec/image")
        
        except UnidentifiedImageError:
            print(f"Unsupported image format for {filename}")
            error_count += 1
        except Exception as e:
            print(f"Error processing {filename}: {e}")
            error_count += 1
    
    average_time = (time.time() - start_time) / max(processed_count, 1)
    print(f"Finished processing. Total Processed: {processed_count}/{total_files}, Errors: {error_count}, Avg. Time: {average_time:.2f} sec/image")
    return f"Processed: {processed_count}/{total_files}, Errors: {error_count}, Remaining: {total_files - processed_count}, Avg. Time: {average_time:.2f} sec"

def process_image(image, input_text, temperature, top_p, top_k, do_sample):
    with torch.no_grad():
        input_by_model = model.build_conversation_input_ids(tokenizer, query=input_text, history=[], images=[image], template_version='base')
        inputs = {
            'input_ids': input_by_model['input_ids'].unsqueeze(0).to(DEVICE),
            'token_type_ids': input_by_model['token_type_ids'].unsqueeze(0).to(DEVICE),
            'attention_mask': input_by_model['attention_mask'].unsqueeze(0).to(DEVICE),
            'images': [[input_by_model['images'][0].to(DEVICE).to(torch_type)]],
        }
        if 'cross_images' in input_by_model and input_by_model['cross_images']:
            inputs['cross_images'] = [[input_by_model['cross_images'][0].to(DEVICE).to(torch_type)]]

        gen_kwargs = {
                "max_length": 2048,
                "temperature": temperature,
                "do_sample": do_sample,
                "top_p": top_p,
                "top_k": top_k
        }
        outputs = model.generate(**inputs, **gen_kwargs)
        outputs = outputs[:, inputs['input_ids'].shape[1]:]
        response = tokenizer.decode(outputs[0])
        return response.split("</s>")[0]

def main():
    gr.close_all()
    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column(scale=4):
                image_prompt = gr.Image(type="filepath", label="Image Prompt", value=None)

                with gr.Row():
                    temperature = gr.Slider(maximum=1, value=0.9, minimum=0, step=0.01, label='Temperature')
                    top_p = gr.Slider(maximum=1, value=0.8, minimum=0, step=0.01, label='Top P')
                    top_k = gr.Slider(maximum=100, value=1, minimum=1, step=1, label='Top K')
                    do_sample_checkbox = gr.Checkbox(label='Enable Sampling', value=True)

                with gr.Row():
                    input_text_examples = gr.components.Textbox(lines=4, label='Prompt Examples', value='Question: Describe this image Answer:\nQuestion: How many people are there? Short answer:', interactive=False)

            with gr.Column(scale=5):
                with gr.Row():
                    input_text = gr.components.Textbox(lines=10, label='Caption Prompt', value='Question: just caption the image with details, colors, items, objects, emotions, art style, drawing style and objects but do not add any description or comment. do not miss any item in the given image Answer:')
                with gr.Row():
                    run_button = gr.Button('Generate', variant='primary')
                with gr.Row():
                    result_text = gr.components.Textbox(lines=4, label='Result Text', placeholder='')

        run_button.click(fn=post, inputs=[input_text, temperature, top_p, top_k, image_prompt, do_sample_checkbox], outputs=[result_text])	

        with gr.Row():
            with gr.Column(scale=4):
                input_folder = gr.Textbox(label="Batch Processing Folder")
                batch_button = gr.Button("Start Batch Processing")
                input_text_batch = gr.Textbox(lines=10, label='Batch Processing Input Text', value='Question: just caption the image with details, colors, items, objects, emotions, art style, drawing style and objects but do not add any description or comment. do not miss any item in the given image Answer:')
                status_label = gr.Label(value="")

        batch_button.click(fn=batch_process, inputs=[input_folder, input_text_batch,temperature, top_p, top_k, do_sample_checkbox], outputs=status_label)

    demo.queue().launch(share=args.share, inbrowser=True)

def post(input_text, temperature, top_p, top_k, image_prompt, do_sample):
    try:
        with torch.no_grad():
            image = Image.open(image_prompt).convert('RGB') if image_prompt is not None else None

            input_by_model = model.build_conversation_input_ids(tokenizer, query=input_text, history=[], images=([image] if image else None), template_version='base')
            inputs = {
                'input_ids': input_by_model['input_ids'].unsqueeze(0).to(DEVICE),
                'token_type_ids': input_by_model['token_type_ids'].unsqueeze(0).to(DEVICE),
                'attention_mask': input_by_model['attention_mask'].unsqueeze(0).to(DEVICE),
                'images': [[input_by_model['images'][0].to(DEVICE).to(torch_type)]],
            }
            if 'cross_images' in input_by_model and input_by_model['cross_images']:
                print('cross images used')
                inputs['cross_images'] = [[input_by_model['cross_images'][0].to(DEVICE).to(torch_type)]]

            gen_kwargs = {
                "max_length": 2048,
                "temperature": temperature,
                "do_sample": do_sample,
                "top_p": top_p,
                "top_k": top_k
            }
            outputs = model.generate(**inputs, **gen_kwargs)
            outputs = outputs[:, inputs['input_ids'].shape[1]:]
            response = tokenizer.decode(outputs[0])
            response = response.split("</s>")[0]
            return response
    except Exception as e:
        return str(e)


if __name__ == "__main__":
    main()
