import requests
import json
from PIL import Image
from io import BytesIO
url = 'https://api-inference.modelscope.cn/v1/images/generations'
payload = {
    'model': 'DiffSynth-Studio/FLUX.1-Kontext-dev-lora-highresfix',#ModelScope Model-Id,required
    'prompt': 'A golden cat'# required
}
headers = {
    'Authorization': 'Bearer <ModelScope-SDK-TOKEN>', # <ModelScope-SDK-TOKEN> 请替换成您的ModelScope SDK Token
    'Content-Type': 'application/json'
}
response = requests.post(url, data=json.dumps(payload, ensure_ascii=False).encode('utf-8'), headers=headers)
response_data = response.json()
image = Image.open(BytesIO(requests.get(response_data['images'][0]['url']).content))
image.save('result_image.jpg')

from mcp.server.fastmcp import FastMCP
from mcp.types import TextContent
import requests
import json
import os
mcp = FastMCP("image_generation_flux_lora")
@mcp.tool()
def generate_image_url_from_text(description : str,
                                 model: str = "MusePublic/489_ckpt_FLUX_1"
                                 ) -> list[TextContent]:
    """Generate an image from the input description using ModelScope API, it returns the image URL.
    Args:
        description: the description of the image to be generated, containing the desired elements and visual features.
        model: the model name to be used for image generation, default is "MusePublic/489_ckpt_FLUX_1".
    """
    url = 'https://api-inference.modelscope.cn/v1/images/generations'
    token = os.environ.get("MODELSCOPE_API_KEY")
    payload = {
        'model': model,  # ModelScope Model-Id, 必填项
        'prompt': description  # 必填项
    }
    headers = {
        'Authorization': f'Bearer {token}',
        'Content-Type': 'application/json'
    }
    try:
        response = requests.post(url,
                                 data=json.dumps(
                                     payload, ensure_ascii=False).encode('utf-8'),
                                 headers=headers)
        response_data = response.json()
        if 'images' in response_data.keys():
            res= response_data['images'][0]['url']
        else:
            res = str(response_data)
    except Exception as e:
        res = "error:" + str(e)
        print(f"Error: {e}")
    return [TextContent(type="text", text=res)]
if __name__ == "__main__":
    mcp.run(transport='stdio')