import requests
import json
import sys
from PIL import Image
import base64

text = sys.argv[1]
image_path = sys.argv[2]

#image = Image.open(image_path)
with open(image_path, "rb") as file:
    image = str(base64.b64encode(file.read()))
print("image:", type(image))
sample = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.###Human: What are the key differences between renewable and non-renewable energy sources?###Assistant: Renewable energy sources are those that can be replenished naturally in a relatively short amount of time, such as solar, wind, hydro, geothermal, and biomass. Non-renewable energy sources, on the other hand, are finite and will eventually be depleted, such as coal, oil, and natural gas. Here are some key differences between renewable and non-renewable energy sources:\n1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable energy sources are finite and will eventually run out.\n2. Environmental impact: Renewable energy sources have a much lower environmental impact than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, and other negative effects.\n3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically have lower operational costs than non-renewable sources.\n4. Reliability: Renewable energy sources are often more reliable and can be used in more remote locations than non-renewable sources.\n5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different situations and needs, while non-renewable sources are more rigid and inflexible.\n6. Sustainability: Renewable energy sources are more sustainable over the long term, while non-renewable sources are not, and their depletion can lead to economic and social instability.\n###Human: "
prompt = sample + text + "\n<image>###Assistant:"

url = "http://task345865.training:6801/worker_generate_stream"
headers = {'User-Agent': 'LLaVA Client'}
pload = {
    'model': 'output_test_checkpoint-6500',
    'prompt': prompt, "images": [image],
    'temperature': 0.5, 'max_new_tokens': 512, 'stop': '###'}

response = requests.post(url,
     headers=headers, json=pload, stream=True, timeout=10)
for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
    if chunk:
        data = json.loads(chunk.decode())
        text = data["text"].split("###")[-1].split(":")[1]
        if text.count("<") == 2 and text.count(">") == 2:
            res = text
        elif text.count("<") == 1 and text.count(">") == 1:
            res = text + text
        else:
            res = "<127><127>"
print(res)