from flask import Flask, request, jsonify
from gradio_client import Client

# Create a Flask object
app = Flask("Generate logo")

client = Client("stabilityai/stable-diffusion-3-medium")

def generate_logo(prompt):
  """Generates a logo using the provided prompt."""

  result = client.predict(
      prompt=prompt,
      negative_prompt="",
      seed=0,
      randomize_seed=True,
      width=1024,
      height=1024,
      guidance_scale=7, 
      num_inference_steps=50,
      api_name="/infer",
  )
  return result

@app.route('/llama', methods=['POST'])
def generate_response():
    try:
        data = request.get_json()
        
        if 'prompt' in data :
            prompt = data['prompt']
            
            response  = generate_logo(prompt)
            print(response, 'response')
            fileUrl = f'https://stabilityai-stable-diffusion-3-medium.hf.space/file={response[0]}'

            return jsonify(fileUrl)

        else:
            return jsonify({"error": "Missing required parameters"}), 400

    except Exception as e:
        return jsonify({"Error": str(e)}), 500 

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)