Spaces:
Sleeping
Sleeping
| from flask import Flask, request, jsonify, send_from_directory | |
| import requests | |
| import io | |
| import random | |
| import os | |
| from PIL import Image | |
| from datetime import datetime | |
| import string | |
| import re | |
| app = Flask(__name__) | |
| API_URL = "https://api-inference.huggingface.co/models/openskyml/midjourney-mini" | |
| API_TOKEN = os.getenv("HF_READ_TOKEN") # it is free | |
| headers = {"Authorization": f"Bearer {API_TOKEN}"} | |
| TEMP_DIR = "temp" | |
| def query(prompt, is_negative=False, steps=1, cfg_scale=6, seed=None): | |
| payload = { | |
| "inputs": prompt, | |
| "is_negative": is_negative, | |
| "steps": steps, | |
| "cfg_scale": cfg_scale, | |
| "seed": seed if seed is not None else random.randint(-1, 2147483647) | |
| } | |
| response = requests.post(API_URL, headers=headers, json=payload) | |
| response.raise_for_status() | |
| image_bytes = response.content | |
| image = Image.open(io.BytesIO(image_bytes)) | |
| return image | |
| def sanitize_filename(filename): | |
| # Replace spaces with underscores | |
| filename = filename.replace(" ", "_") | |
| # Remove special characters | |
| sanitized_filename = re.sub(r'[^\w\s-]', '', filename) | |
| sanitized_filename = re.sub(r'[-\s]+', '-', sanitized_filename) | |
| return sanitized_filename | |
| def save_image(image, prompt): | |
| timestamp = datetime.now().strftime("%Y%m%d%H%M%S") | |
| sanitized_prompt = sanitize_filename(prompt) | |
| filename = f"{sanitized_prompt}-{random.randint(1, 100000)}-{timestamp}.png" | |
| filepath = os.path.join(TEMP_DIR, filename) | |
| image.save(filepath, format='PNG') | |
| return filename | |
| def generate(): | |
| try: | |
| data = request.get_json() | |
| prompt = data["prompt"] | |
| negative_prompt = data.get("negative_prompt", "") | |
| is_negative = True if negative_prompt else False | |
| image = query(prompt, is_negative=is_negative) | |
| filename = save_image(image, prompt) | |
| response = { | |
| "success": True, | |
| "image": f"https://mrdonstuff-dalle-3-xl-api.hf.space/temp/{filename}", | |
| "filename": filename | |
| } | |
| except Exception as e: | |
| response = { | |
| "success": False, | |
| "error": str(e) | |
| } | |
| return jsonify(response) | |
| def show_image(filename): | |
| return send_from_directory(TEMP_DIR, filename) | |
| if __name__ == "__main__": | |
| if not os.path.exists(TEMP_DIR): | |
| os.makedirs(TEMP_DIR) | |
| app.run(debug=True, host="0.0.0.0", port=7860) |