Created
June 12, 2025 09:41
-
-
Save MikyPo/d29642919fc79e6568a916e84f3dc5b4 to your computer and use it in GitHub Desktop.
kandinsky_fusionbrain_api
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import base64 | |
| import json | |
| import os | |
| import requests | |
| import time | |
| from io import BytesIO | |
| from dotenv import load_dotenv | |
| from IPython.display import display, Image | |
| load_dotenv('.env') | |
| YOUR_API_KEY = os.getenv('YOUR_API_KEY') | |
| YOUR_SECRET_KEY = os.getenv('YOUR_SECRET_KEY') | |
| API_URL = 'https://api-key.fusionbrain.ai/' | |
| class FusionBrainAPI: | |
| def __init__(self, url, api_key, secret_key): | |
| self.URL = url | |
| self.AUTH_HEADERS = { | |
| 'X-Key': f'Key {api_key}', | |
| 'X-Secret': f'Secret {secret_key}', | |
| } | |
| def get_pipeline(self): | |
| """Получаем ID доступного пайплайна (Kandinsky)""" | |
| try: | |
| response = requests.get(self.URL + 'key/api/v1/pipelines', headers=self.AUTH_HEADERS) | |
| response.raise_for_status() | |
| data = response.json() | |
| if not data: | |
| raise Exception("Нет доступных пайплайнов. Проверьте API ключи.") | |
| for pipeline in data: | |
| if "kandinsky" in pipeline.get("name", "").lower(): | |
| return pipeline['id'] | |
| raise Exception("Пайплайн Kandinsky не найден.") | |
| except Exception as e: | |
| raise Exception(f"Ошибка получения пайплайна: {str(e)}") | |
| def generate(self, prompt, style="DEFAULT", negative_prompt="", width=512, height=512, num_images=1): | |
| """Генерация изображения""" | |
| try: | |
| pipeline_id = self.get_pipeline() | |
| params = { | |
| "type": "GENERATE", | |
| "style": style, | |
| "width": width, | |
| "height": height, | |
| "numImages": num_images, | |
| "negativePromptDecoder": negative_prompt, | |
| "generateParams": { | |
| "query": prompt | |
| } | |
| } | |
| files = { | |
| 'pipeline_id': (None, pipeline_id), | |
| 'params': (None, json.dumps(params), 'application/json') | |
| } | |
| response = requests.post( | |
| self.URL + 'key/api/v1/pipeline/run', | |
| headers=self.AUTH_HEADERS, | |
| files=files | |
| ) | |
| response.raise_for_status() | |
| return response.json()['uuid'] | |
| except Exception as e: | |
| raise Exception(f"Ошибка генерации: {str(e)}") | |
| def check_generation(self, request_id, attempts=10, delay=5): | |
| """Проверка статуса генерации""" | |
| try: | |
| while attempts > 0: | |
| response = requests.get( | |
| self.URL + 'key/api/v1/pipeline/status/' + request_id, | |
| headers=self.AUTH_HEADERS | |
| ) | |
| response.raise_for_status() | |
| data = response.json() | |
| if data['status'] == 'DONE': | |
| return data['result']['files'] | |
| elif data['status'] == 'FAILED': | |
| return None | |
| attempts -= 1 | |
| time.sleep(delay) | |
| return None | |
| except Exception as e: | |
| raise Exception(f"Ошибка проверки статуса: {str(e)}") | |
| def get_image_bytes(self, base64_string): | |
| """Конвертация Base64 в bytes""" | |
| if base64_string.startswith('data:image'): | |
| base64_string = base64_string.split(',')[1] | |
| return base64.b64decode(base64_string) | |
| def generate_and_show(self, prompt, **kwargs): | |
| """Генерация + отображение в Jupyter""" | |
| try: | |
| print("⌛ Генерация начата...") | |
| print() | |
| uuid = self.generate(prompt, **kwargs) | |
| files = self.check_generation(uuid) | |
| if files: | |
| image_data = self.get_image_bytes(files[0]) | |
| display(Image(image_data)) | |
| print() | |
| print("✅ Готово!") | |
| return image_data # Возвращаем данные изображения | |
| else: | |
| print("❌ Ошибка генерации") | |
| return None | |
| except Exception as e: | |
| print(f"❌ Ошибка: {str(e)}") | |
| return None | |
| api = FusionBrainAPI(API_URL, YOUR_API_KEY, YOUR_SECRET_KEY) | |
| # Стиль из списка https://cdn.fusionbrain.ai/static/styles/key | |
| # Возможные значения: | |
| # "DEFAULT" - "No style" | |
| # "KANDINSKY" - "Kandinsky" | |
| # "UHD" - "Detailed photo" | |
| # "ANIME" - "Anime" | |
| PROMPT = "киборг делает пирог" | |
| STYLE = "KANDINSKY" | |
| NEGATIVE_PROMPT = "размытость, низкое качество" | |
| image_data = api.generate_and_show( | |
| prompt=PROMPT, | |
| style=STYLE, | |
| negative_prompt=NEGATIVE_PROMPT, | |
| # Рекомендуется использовать значения кратные 64 с каждой из сторон | |
| width=1024, | |
| height=768 | |
| ) | |
| # Сохранение файла | |
| if image_data: | |
| with open(f"result_{STYLE}_{PROMPT}.jpg", "wb") as f: | |
| f.write(image_data) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment