Create app.py

app.py

@@ -0,0 +1,75 @@
+import gradio as gr
+import torch
+import json
+from io import BytesIO
+from PIL import Image, ImageOps
+from IPython.display import display, Markdown
+from transformers import AutoModelForCausalLM, LlamaTokenizer
+from accelerate import init_empty_weights, infer_auto_device_map, load_checkpoint_and_dispatch
+
+# Initialize tokenizer and model
+tokenizer = LlamaTokenizer.from_pretrained('lmsys/vicuna-7b-v1.5')
+# tokenizer = LlamaTokenizer.from_pretrained('vicuna-7b-v1.5')
+model = AutoModelForCausalLM.from_pretrained(
+        'THUDM/cogvlm-chat-hf',
+        load_in_4bit=True,
+        trust_remote_code=True,
+        device_map="auto"
+    ).eval()
+
+def generate_description(image, query, top_p, top_k, output_length, temperature):
+    # Use the uploaded image (PIL format)
+    display_size = (224, 224)
+    image = image.resize(display_size, Image.LANCZOS)
+
+    # Build the conversation input
+    inputs = model.build_conversation_input_ids(tokenizer, query=query, history=[], images=[image])
+
+    # Prepare the inputs dictionary for model.generate()
+    inputs = {
+        'input_ids': inputs['input_ids'].unsqueeze(0).to('cuda'),
+        'token_type_ids': inputs['token_type_ids'].unsqueeze(0).to('cuda'),
+        'attention_mask': inputs['attention_mask'].unsqueeze(0).to('cuda'),
+        'images': [[inputs['images'][0].to('cuda').to(torch.float16)]],
+    }
+
+    # Set the generation kwargs with user-defined values
+    gen_kwargs = {
+        "max_length": output_length,
+        "do_sample": True,  # Enable sampling to use top_p, top_k, and temperature
+        "top_p": top_p,
+        "top_k": top_k,
+        "temperature": temperature
+    }
+
+    # Generate the description
+    with torch.no_grad():
+        outputs = model.generate(**inputs, **gen_kwargs)
+        description = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    return description
+
+with gr.Blocks() as app:
+    gr.Markdown("# Visual Product DNA - Image to Attribute Extractor")
+    
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(label="Upload Image", type="pil", height=500) 
+            gr.skip 
+            query_input = gr.Textbox(label="Enter your prompt", value="Capture all attributes as JSON", lines=4)  
+            
+        with gr.Column():
+            top_p_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.1, label="Creativity (top_p)")
+            top_k_slider = gr.Slider(minimum=0, maximum=100, step=1, value=100, label="Coherence (top_k)")
+            output_length_slider = gr.Slider(minimum=1, maximum=4096, step=1, value=2048, label="Output Length")
+            temperature_slider = gr.Slider(minimum=0.1, maximum=2.0, step=0.01, value=0.1, label="Temperature")
+            submit_button = gr.Button("Extract Attributes")
+            description_output = gr.Textbox(label="Generated JSON", lines=12)  
+            
+    submit_button.click(
+        fn=generate_description, 
+        inputs=[image_input, query_input, top_p_slider, top_k_slider, output_length_slider, temperature_slider], 
+        outputs=description_output
+    )
+
+app.launch(share=True, input = False)