imagine.py

import aiohttp
import asyncio

@bot.hybrid_command(name="imagine", description="Generate image using an endpoint")
async def images(ctx, *, prompt):
    api_endpoint = "https://api.salad.com/api/beta/organizations/erisws/inference-endpoints/stable-diffusion-2-1"
    headers = {
        "Content-Type": "application/json",
        "Salad-Api-Key": "MY-API-KEY"
    }
    data = {
        "modelInputs": {
            "prompt": prompt,
            "num_inference_steps": 75,
            "guidance_scale": 20,
            "width": 1024,
            "height": 1024,
        },
        "callInputs": {
            "PIPELINE": "StableDiffusionPipeline",
            "SCHEDULER": "LMSDiscreteScheduler",
            "safety_checker": "true",
        }
    }
    try:
        temp_message = await ctx.send("Generating image...")
        
        async with aiohttp.ClientSession() as session:
            async with session.post(api_endpoint, headers=headers, json=data) as response:
                response_json = await response.json()
                request_id = response_json["requestId"]
                result_url = f"https://api.salad.com/api/beta/organizations/erisws/inference-endpoints/stable-diffusion-2-1/{request_id}"
                image_exists = False
                while not image_exists:
                    await asyncio.sleep(10)
                    async with session.get(result_url) as result_response:
                        result_json = await result_response.json()
                        if 'modelOutputs' in result_json and 'image_base64' in result_json['modelOutputs']:
                            image_base64 = result_json["modelOutputs"]["image_base64"]
                            image_exists = True
                        else:
                            print("Image value not present yet, will wait 10 seconds...")
                
                image_bytes = base64.b64decode(image_base64)
                image_name = f"{prompt}.png"
                with open(image_name, "wb") as file:
                    file.write(image_bytes)
                    
                await ctx.send(
                    f"Prompt by {ctx.author.mention} : `{prompt}`",
                    file=discord.File(image_name, filename=f"{image_name}")
                )
                
                await temp_message.edit(content="Finished Image Generation")
                
                os.remove(image_name)
    
    except aiohttp.ClientError as e:
        await temp_message.edit(content=f"An error occurred while sending the request: {str(e)}")
    
    except Exception as e:
        await temp_message.edit(content=f"An error occurred: {str(e)}")