Files

568 lines
27 KiB
Python
Raw Permalink Normal View History

2026-05-31 18:45:24 +08:00
import sys
import json
import os
import random
import urllib.request
import urllib.parse
import websocket
import requests
import uuid
import io
import logging
import re
import copy
from PIL import Image
from PIL.PngImagePlugin import PngInfo
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
QHBoxLayout, QTextEdit, QSpinBox, QComboBox,
QPushButton, QLabel, QFileDialog, QInputDialog, QMessageBox,
QCheckBox, QProgressBar, QFrame, QGroupBox, QSizePolicy)
from PyQt5.QtGui import QPixmap, QImage, QFont, QIcon
from PyQt5.QtCore import Qt, QThread, pyqtSignal
# ================= ЛОГИРОВАНИЕ =================
logger = logging.getLogger("ComfyUIRemote")
logger.setLevel(logging.DEBUG)
formatter = logging.Formatter('%(asctime)s | %(levelname)-8s | %(message)s')
file_handler = logging.FileHandler("app_detailed_log.txt", encoding='utf-8')
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
# ================= КОНФИГУРАЦИЯ =================
CONFIG_FILE = "settings.json"
DEFAULT_CONFIG = {
"comfyui": {"server_address": "192.168.1.118:8188"},
"lmstudio": {
"base_url": "http://192.168.1.118:1234/v1",
"api_key": "lm-studio",
"models": ["luna-ai-llama2"]
},
"cloud_ai": {
"base_url": "https://neuroapi.host/v1",
"api_key": "ВАШ_КЛЮЧ_ЗДЕСЬ",
"models": ["gemini-2.5-flash", "deepseek-v3.2", "grok-4-fast-reasoning"]
},
"ollama": {
"base_url": "http://127.0.0.1:11434/v1"
}
}
# ================= WORKFLOWS =================
# 1. Базовый Text-to-Image (стандартный изначальный воркфлоу)
WORKFLOW_T2I = {
"39": {"inputs": {"filename_prefix": "ComfyUI", "images": ["35:8", 0]}, "class_type": "SaveImage"},
"35:30": {"inputs": {"clip_name": "qwen_3_4b.safetensors", "type": "lumina2", "device": "default"}, "class_type": "CLIPLoader"},
"35:29": {"inputs": {"vae_name": "ae.safetensors"}, "class_type": "VAELoader"},
"35:28": {"inputs": {"unet_name": "SwarmUI_Z-Image-Turbo-FP8Mix.safetensors", "weight_dtype": "default"}, "class_type": "UNETLoader"},
"35:27": {"inputs": {"text": "", "clip": ["35:30", 0]}, "class_type": "CLIPTextEncode"},
"35:13": {"inputs": {"width": 1024, "height": 1024, "batch_size": 1}, "class_type": "EmptySD3LatentImage"},
"35:3": {"inputs": {"seed": 0, "steps": 4, "cfg": 1, "sampler_name": "res_multistep", "scheduler": "simple", "denoise": 1, "model": ["35:11", 0], "positive": ["35:27", 0], "negative": ["35:33", 0], "latent_image": ["35:13", 0]}, "class_type": "KSampler"},
"35:11": {"inputs": {"shift": 3, "model": ["35:28", 0]}, "class_type": "ModelSamplingAuraFlow"},
"35:33": {"inputs": {"conditioning": ["35:27", 0]}, "class_type": "ConditioningZeroOut"},
"35:8": {"inputs": {"samples": ["35:3", 0], "vae": ["35:29", 0]}, "class_type": "VAEDecode"}
}
# 2. Image Editing (Qwen-Image-Edit)
WORKFLOW_I2I = {
"60": {"inputs": {"filename_prefix": "ComfyUI", "images": ["102:8", 0]}, "class_type": "SaveImage"},
"78": {"inputs": {"image": "image_qwen_image_edit_input_image.png"}, "class_type": "LoadImage"},
"93": {"inputs": {"upscale_method": "lanczos", "megapixels": 1.5, "resolution_steps": 1, "image": ["78", 0]}, "class_type": "ImageScaleToTotalPixels"},
"102:39": {"inputs": {"vae_name": "qwen_image_vae.safetensors"}, "class_type": "VAELoader"},
"102:77": {"inputs": {"prompt": "", "clip": ["102:38", 0], "vae": ["102:39", 0], "image": ["78", 0]}, "class_type": "TextEncodeQwenImageEdit"},
"102:75": {"inputs": {"strength": 1, "model": ["102:66", 0]}, "class_type": "CFGNorm"},
"102:66": {"inputs": {"shift": 3, "model": ["102:108", 0]}, "class_type": "ModelSamplingAuraFlow"},
"102:8": {"inputs": {"samples": ["102:3", 0], "vae": ["102:39", 0]}, "class_type": "VAEDecode"},
"102:38": {"inputs": {"clip_name": "qwen_2.5_vl_7b_fp8_scaled.safetensors", "type": "qwen_image", "device": "default"}, "class_type": "CLIPLoader"},
"102:76": {"inputs": {"prompt": "", "clip": ["102:38", 0], "vae": ["102:39", 0], "image": ["78", 0]}, "class_type": "TextEncodeQwenImageEdit"},
"102:88": {"inputs": {"pixels": ["78", 0], "vae": ["102:39", 0]}, "class_type": "VAEEncode"},
"102:89": {"inputs": {"lora_name": "Qwen-Image-Edit-Lightning-4steps-V1.0-bf16.safetensors", "strength_model": 1, "model": ["102:37", 0]}, "class_type": "LoraLoaderModelOnly"},
"102:37": {"inputs": {"unet_name": "qwen_image_edit_fp8_e4m3fn.safetensors", "weight_dtype": "default"}, "class_type": "UNETLoader"},
"102:105": {"inputs": {"value": 1}, "class_type": "PrimitiveFloat"},
"102:106": {"inputs": {"value": 20}, "class_type": "PrimitiveInt"},
"102:103": {"inputs": {"value": 4}, "class_type": "PrimitiveInt"},
"102:107": {"inputs": {"value": 2.5}, "class_type": "PrimitiveFloat"},
"102:111": {"inputs": {"value": False}, "class_type": "PrimitiveBoolean"},
"102:3": {"inputs": {"seed": 0, "steps": ["102:110", 0], "cfg": ["102:109", 0], "sampler_name": "euler", "scheduler": "simple", "denoise": 1, "model": ["102:75", 0], "positive": ["102:76", 0], "negative": ["102:77", 0], "latent_image": ["102:88", 0]}, "class_type": "KSampler"},
"102:109": {"inputs": {"switch": ["102:111", 0], "on_false": ["102:107", 0], "on_true": ["102:105", 0]}, "class_type": "ComfySwitchNode"},
"102:110": {"inputs": {"switch": ["102:111", 0], "on_false": ["102:106", 0], "on_true": ["102:103", 0]}, "class_type": "ComfySwitchNode"},
"102:108": {"inputs": {"switch": ["102:111", 0], "on_false": ["102:37", 0], "on_true": ["102:89", 0]}, "class_type": "ComfySwitchNode"}
}
# ================= СТИЛИ (QSS) =================
DARK_STYLE = """
QMainWindow { background-color: #1e1e2e; }
QGroupBox { color: #cdd6f4; font-weight: bold; border: 1px solid #45475a; margin-top: 15px; border-radius: 8px; padding-top: 20px; }
QLabel { color: #bac2de; font-size: 13px; }
QTextEdit { background-color: #313244; color: #cdd6f4; border: 1px solid #45475a; border-radius: 5px; padding: 10px; selection-background-color: #585b70; }
QSpinBox, QComboBox { background-color: #313244; color: #cdd6f4; border: 1px solid #45475a; border-radius: 4px; padding: 5px; min-height: 30px; }
QPushButton { background-color: #45475a; color: #cdd6f4; border-radius: 6px; padding: 10px; font-weight: bold; }
QPushButton:hover { background-color: #585b70; }
QPushButton#start_btn { background-color: #a6e3a1; color: #11111b; font-size: 14px; }
QPushButton#start_btn:hover { background-color: #94e2d5; }
QPushButton#save_btn { background-color: #f38ba8; color: #11111b; }
QProgressBar { border: 1px solid #45475a; border-radius: 5px; text-align: center; color: white; height: 15px; }
QProgressBar::chunk { background-color: #89b4fa; }
QCheckBox { color: #bac2de; font-size: 13px; spacing: 8px; }
"""
# ================= WORKERS =================
class ComfyUIWorker(QThread):
finished_signal = pyqtSignal(QPixmap, bytes)
error_signal = pyqtSignal(str)
progress_signal = pyqtSignal(int, int)
def __init__(self, workflow, server_address, upload_image_path=None, image_node_id=None):
super().__init__()
self.workflow = workflow
self.server_address = server_address
self.client_id = str(uuid.uuid4())
self.upload_image_path = upload_image_path
self.image_node_id = image_node_id
def run(self):
ws = None
try:
# 1. Если задан путь к картинке (I2I Workflow) -> загружаем ее на сервер
if self.upload_image_path and self.image_node_id:
try:
with open(self.upload_image_path, 'rb') as f:
upload_url = f"http://{self.server_address}/upload/image"
files = {'image': f}
req = requests.post(upload_url, files=files)
if req.status_code == 200:
# Получаем реальное имя файла на сервере
filename = req.json()['name']
self.workflow[self.image_node_id]["inputs"]["image"] = filename
else:
self.error_signal.emit("Failed to upload image to ComfyUI.")
return
except Exception as e:
self.error_signal.emit(f"Image upload error: {str(e)}")
return
# 2. Подключаемся к WebSocket и отправляем workflow
ws_url = f"ws://{self.server_address}/ws?clientId={self.client_id}"
ws = websocket.WebSocket()
ws.connect(ws_url)
prompt_url = f"http://{self.server_address}/prompt"
data = json.dumps({"prompt": self.workflow, "client_id": self.client_id}).encode('utf-8')
req = urllib.request.Request(prompt_url, data=data)
with urllib.request.urlopen(req) as response:
prompt_id = json.loads(response.read())['prompt_id']
# Ждем выполнения
while True:
out = ws.recv()
if isinstance(out, str):
message = json.loads(out)
if message['type'] == 'progress':
self.progress_signal.emit(message['data']['value'], message['data']['max'])
if message['type'] == 'executing' and message['data']['node'] is None:
if message['data']['prompt_id'] == prompt_id: break
# Получаем результат
with urllib.request.urlopen(f"http://{self.server_address}/history/{prompt_id}") as response:
history = json.loads(response.read())[prompt_id]
for node_id in history['outputs']:
node_output = history['outputs'][node_id]
if 'images' in node_output:
img_data = node_output['images'][0]
view_url = f"http://{self.server_address}/view?filename={img_data['filename']}&subfolder={img_data['subfolder']}&type={img_data['type']}"
with urllib.request.urlopen(view_url) as img_res:
img_bytes = img_res.read()
self.finished_signal.emit(QPixmap.fromImage(QImage.fromData(img_bytes)), img_bytes)
return
self.error_signal.emit("Generation finished, but no image found in outputs.")
except Exception as e:
self.error_signal.emit(str(e))
finally:
if ws: ws.close()
class LLMWorker(QThread):
finished_signal = pyqtSignal(str)
error_signal = pyqtSignal(str)
def __init__(self, base_url, api_key, model, prompt, provider="default"):
super().__init__()
self.base_url = base_url.rstrip('/')
self.api_key = api_key
self.model = model
self.prompt = prompt
self.provider = provider
def run(self):
try:
headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
if self.provider == "ollama":
system_instruction = (
"You are a professional image prompt engineer. "
"Task: Translate the user's input into ENGLISH (if it's not already) and expand it into a highly detailed visual prompt. "
"CRITICAL RULES: "
"1. MUST BE IN ENGLISH ONLY. Do NOT output Russian words. "
"2. NO REASONING, NO THOUGHTS. Do not explain your process. "
"3. OUTPUT EXACTLY THE RAW PROMPT. No introductions like 'Here is...' or 'Sure'. "
)
temperature = 0.6
else:
system_instruction = "Improve prompt for AI image generation. Expanded English description only."
temperature = 0.7
payload = {
"model": self.model,
"messages": [
{"role": "system", "content": system_instruction},
{"role": "user", "content": self.prompt}
],
"temperature": temperature
}
response = requests.post(f"{self.base_url}/chat/completions", headers=headers, json=payload, timeout=180)
if response.status_code == 200:
result_text = response.json()['choices'][0]['message']['content'].strip()
result_text = re.sub(r'<think>.*?</think>', '', result_text, flags=re.DOTALL).strip()
if self.provider == "ollama":
if result_text.startswith('"') and result_text.endswith('"'):
result_text = result_text[1:-1].strip()
self.finished_signal.emit(result_text)
else:
self.error_signal.emit(f"Error {response.status_code}: {response.text}")
except requests.exceptions.Timeout:
self.error_signal.emit("Тайм-аут: Сервер слишком долго загружал модель или генерировал ответ. Попробуйте еще раз.")
except Exception as e:
self.error_signal.emit(str(e))
# ================= ГЛАВНОЕ ОКНО =================
class MainWindow(QMainWindow):
def __init__(self):
super().__init__()
self.setWindowTitle("ComfyUI Studio v3")
self.resize(1280, 850)
self.setStyleSheet(DARK_STYLE)
self.load_config()
self.current_bytes = None
self.selected_image_path = None
self.init_ui()
def load_config(self):
if not os.path.exists(CONFIG_FILE):
with open(CONFIG_FILE, 'w', encoding='utf-8') as f: json.dump(DEFAULT_CONFIG, f, indent=4)
self.config = DEFAULT_CONFIG
else:
try:
with open(CONFIG_FILE, 'r', encoding='utf-8') as f: self.config = json.load(f)
except: self.config = DEFAULT_CONFIG
def init_ui(self):
central_widget = QWidget()
self.setCentralWidget(central_widget)
main_layout = QHBoxLayout(central_widget)
main_layout.setSpacing(20)
main_layout.setContentsMargins(15, 15, 15, 15)
left_panel = QFrame()
left_panel.setMinimumWidth(550)
left_layout = QVBoxLayout(left_panel)
left_layout.setContentsMargins(0, 0, 0, 0)
# Выбор Workflow
wf_group = QGroupBox("Workflow Selector")
wfg_layout = QVBoxLayout()
self.workflow_cb = QComboBox()
self.workflow_cb.addItems([
"Text-to-Image (Default Turbo)",
"Image Editing (Qwen I2I)"
])
self.workflow_cb.currentIndexChanged.connect(self.on_workflow_changed)
wfg_layout.addWidget(self.workflow_cb)
wf_group.setLayout(wfg_layout)
left_layout.addWidget(wf_group)
# Prompt Input
prompt_group = QGroupBox("Prompt Input")
pg_layout = QVBoxLayout()
self.prompt_edit = QTextEdit()
self.prompt_edit.setPlaceholderText("Describe your image idea here...")
self.prompt_edit.setMinimumHeight(150)
pg_layout.addWidget(self.prompt_edit)
self.neg_check = QCheckBox("Use Negative Prompt")
self.neg_check.toggled.connect(lambda c: self.neg_edit.setVisible(c))
pg_layout.addWidget(self.neg_check)
self.neg_edit = QTextEdit()
self.neg_edit.setPlaceholderText("Describe what to exclude...")
self.neg_edit.setMinimumHeight(80)
self.neg_edit.setVisible(False)
pg_layout.addWidget(self.neg_edit)
prompt_group.setLayout(pg_layout)
left_layout.addWidget(prompt_group, 3)
# Блок Изображения (Только для Image Editing)
self.image_picker_group = QGroupBox("Input Image Settings")
ip_layout = QVBoxLayout()
btn_img_layout = QHBoxLayout()
self.btn_select_image = QPushButton("Select Local Image")
self.btn_select_image.clicked.connect(self.select_input_image)
self.lbl_input_image = QLabel("No image selected")
self.lbl_input_image.setWordWrap(True)
btn_img_layout.addWidget(self.btn_select_image)
btn_img_layout.addWidget(self.lbl_input_image)
self.fast_edit_cb = QCheckBox("Enable 4-Steps LoRA (Fast Edit)")
self.fast_edit_cb.setChecked(True)
ip_layout.addLayout(btn_img_layout)
ip_layout.addWidget(self.fast_edit_cb)
self.image_picker_group.setLayout(ip_layout)
self.image_picker_group.setVisible(False) # Скрыто по умолчанию
left_layout.addWidget(self.image_picker_group)
# Settings Group
settings_group = QGroupBox("Core Parameters")
sg_layout = QVBoxLayout()
row1 = QHBoxLayout()
self.step_s = QSpinBox(); self.step_s.setRange(1, 100); self.step_s.setValue(4)
row1.addWidget(QLabel("Steps:")); row1.addWidget(self.step_s)
self.seed_s = QSpinBox(); self.seed_s.setRange(0, 2147483647); self.seed_s.setValue(random.randint(0, 1000000))
row1.addWidget(QLabel("Seed:")); row1.addWidget(self.seed_s)
sg_layout.addLayout(row1)
row2 = QHBoxLayout()
self.size_cb = QComboBox()
self.size_cb.addItems(["1024x1024", "832x1216", "1216x832", "1920x1080"])
row2.addWidget(QLabel("Resolution:")); row2.addWidget(self.size_cb)
self.font_cb = QComboBox()
self.font_cb.addItems(["10", "12", "14", "16", "18", "20", "24", "26"]); self.font_cb.setCurrentText("14")
self.font_cb.currentIndexChanged.connect(self.update_font)
row2.addWidget(QLabel("Text Size:")); row2.addWidget(self.font_cb)
sg_layout.addLayout(row2)
self.vae_cb = QComboBox()
self.vae_cb.addItem("ae.safetensors")
sg_layout.addWidget(QLabel("VAE Checkpoint (T2I):"))
sg_layout.addWidget(self.vae_cb)
settings_group.setLayout(sg_layout)
left_layout.addWidget(settings_group, 1)
# AI Tools Group
ai_group = QGroupBox("AI Prompt Assistance")
ag_layout = QHBoxLayout()
self.btn_lms = QPushButton("LM Studio")
self.btn_lms.clicked.connect(self.use_lms)
self.btn_ollama = QPushButton("Ollama")
self.btn_ollama.setStyleSheet("color: #fab387;")
self.btn_ollama.clicked.connect(self.use_ollama)
self.btn_ai = QPushButton("NeuroAPI")
self.btn_ai.setStyleSheet("color: #ca9ee6;")
self.btn_ai.clicked.connect(self.use_ai)
ag_layout.addWidget(self.btn_lms)
ag_layout.addWidget(self.btn_ollama)
ag_layout.addWidget(self.btn_ai)
ai_group.setLayout(ag_layout)
left_layout.addWidget(ai_group)
main_layout.addWidget(left_panel, 2)
# ПРАВАЯ ПАНЕЛЬ (ПРОСМОТР)
right_panel = QFrame()
right_layout = QVBoxLayout(right_panel)
right_layout.setContentsMargins(0, 0, 0, 0)
self.img_lbl = QLabel("Image Viewer")
self.img_lbl.setAlignment(Qt.AlignCenter)
self.img_lbl.setStyleSheet("background-color: #11111b; border-radius: 10px; border: 2px solid #313244;")
self.img_lbl.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
right_layout.addWidget(self.img_lbl)
self.pbar = QProgressBar()
right_layout.addWidget(self.pbar)
controls = QHBoxLayout()
self.btn_start = QPushButton("▶ START GENERATION")
self.btn_start.setObjectName("start_btn")
self.btn_start.setFixedHeight(55)
self.btn_start.clicked.connect(self.run_comfy)
btn_upd = QPushButton("🔄")
btn_upd.setFixedWidth(60); btn_upd.setFixedHeight(55)
btn_upd.clicked.connect(self.refresh)
btn_sav = QPushButton("💾 SAVE IMAGE")
btn_sav.setObjectName("save_btn")
btn_sav.setFixedHeight(55)
btn_sav.clicked.connect(self.save_img)
controls.addWidget(self.btn_start, 3)
controls.addWidget(btn_upd, 1)
controls.addWidget(btn_sav, 2)
right_layout.addLayout(controls)
main_layout.addWidget(right_panel, 3)
self.update_font()
self.fetch_vaes_from_comfy()
def update_font(self):
f = QFont(); f.setPointSize(int(self.font_cb.currentText()))
self.prompt_edit.setFont(f); self.neg_edit.setFont(f)
def fetch_vaes_from_comfy(self):
server = self.config['comfyui']['server_address']
try:
resp = requests.get(f"http://{server}/object_info/VAELoader", timeout=2)
if resp.status_code == 200:
vaes = resp.json()['VAELoader']['input']['required']['vae_name'][0]
self.vae_cb.clear(); self.vae_cb.addItems(vaes)
except: pass
def on_workflow_changed(self):
idx = self.workflow_cb.currentIndex()
if idx == 1: # Qwen Image Edit (I2I)
self.image_picker_group.setVisible(True)
self.step_s.setEnabled(False) # Регулируется через Fast Mode
self.size_cb.setEnabled(False) # Регулируется через исходное изображение
self.vae_cb.setEnabled(False) # Воркфлоу требует специфичный qwen_image_vae
else: # Text-to-Image Default
self.image_picker_group.setVisible(False)
self.step_s.setEnabled(True)
self.size_cb.setEnabled(True)
self.vae_cb.setEnabled(True)
def select_input_image(self):
path, _ = QFileDialog.getOpenFileName(self, "Select Input Image", "", "Images (*.png *.jpg *.jpeg *.webp)")
if path:
self.selected_image_path = path
self.lbl_input_image.setText(os.path.basename(path))
def run_comfy(self):
txt = self.prompt_edit.toPlainText().strip()
if not txt:
QMessageBox.warning(self, "Warning", "Prompt cannot be empty!")
return
self.btn_start.setEnabled(False)
self.btn_start.setText("G E N E R A T I N G . . .")
idx = self.workflow_cb.currentIndex()
upload_path = None
image_node = None
if idx == 0:
# --- Text to Image ---
wf = copy.deepcopy(WORKFLOW_T2I)
wf["35:27"]["inputs"]["text"] = txt
wf["35:3"]["inputs"]["steps"] = self.step_s.value()
wf["35:3"]["inputs"]["seed"] = self.seed_s.value()
w, h = map(int, self.size_cb.currentText().split('x'))
wf["35:13"]["inputs"]["width"] = w
wf["35:13"]["inputs"]["height"] = h
if self.vae_cb.currentText():
wf["35:29"]["inputs"]["vae_name"] = self.vae_cb.currentText()
elif idx == 1:
# --- Image to Image / Qwen Edit ---
if not self.selected_image_path:
QMessageBox.warning(self, "Warning", "Please select an input image first!")
self.btn_start.setEnabled(True)
self.btn_start.setText("▶ START GENERATION")
return
wf = copy.deepcopy(WORKFLOW_I2I)
wf["102:76"]["inputs"]["prompt"] = txt
neg_txt = self.neg_edit.toPlainText().strip() if self.neg_check.isChecked() else ""
wf["102:77"]["inputs"]["prompt"] = neg_txt
wf["102:3"]["inputs"]["seed"] = self.seed_s.value()
# Включаем или выключаем 4-steps LoRA (True = 4 Steps/CFG 1, False = 20 Steps/CFG 2.5)
wf["102:111"]["inputs"]["value"] = self.fast_edit_cb.isChecked()
upload_path = self.selected_image_path
image_node = "78" # Узел LoadImage в Qwen-воркфлоу
# Передаем параметры в Worker
self.worker = ComfyUIWorker(wf, self.config['comfyui']['server_address'], upload_path, image_node)
self.worker.finished_signal.connect(self.on_comfy_finished)
self.worker.error_signal.connect(self.on_comfy_error)
self.worker.progress_signal.connect(lambda v, m: self.pbar.setValue(int(v/m*100) if m > 0 else 0))
self.worker.start()
def on_comfy_finished(self, pix, raw):
self.current_bytes = raw
self.img_lbl.setPixmap(pix.scaled(self.img_lbl.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))
self.btn_start.setEnabled(True)
self.btn_start.setText("▶ START GENERATION")
def on_comfy_error(self, err):
QMessageBox.critical(self, "Generation Error", err)
self.btn_start.setEnabled(True)
self.btn_start.setText("▶ START GENERATION")
def refresh(self):
self.seed_s.setValue(random.randint(0, 2147483647))
self.run_comfy()
def save_img(self):
if not self.current_bytes: return
path, _ = QFileDialog.getSaveFileName(self, "Save Image", "generated_art.png", "*.png")
if path:
img = Image.open(io.BytesIO(self.current_bytes))
img.save(path)
def use_lms(self):
models = self.config['lmstudio'].get('models', [])
m, ok = QInputDialog.getItem(self, "LM Studio", "Select Model:", models, 0, False)
if ok: self._run_llm(self.config['lmstudio']['base_url'], self.config['lmstudio']['api_key'], m, "lmstudio")
def use_ollama(self):
url = self.config.get('ollama', {}).get('base_url', 'http://127.0.0.1:11434/v1')
try:
resp = requests.get(f"{url}/models", timeout=3)
if resp.status_code == 200:
models = [m['id'] for m in resp.json().get('data', [])]
else:
raise Exception(f"HTTP Error {resp.status_code}")
except Exception as e:
QMessageBox.warning(self, "Ollama Error", f"Failed to fetch models from Ollama.\nEnsure Ollama is running.\nError: {e}")
return
if not models:
QMessageBox.warning(self, "Ollama Info", "No models found. Please pull a model first.")
return
m, ok = QInputDialog.getItem(self, "Ollama", "Select Model:", models, 0, False)
if ok: self._run_llm(url, "ollama", m, "ollama")
def use_ai(self):
models = self.config['cloud_ai'].get('models', [])
m, ok = QInputDialog.getItem(self, "NeuroAPI", "Select Model:", models, 0, False)
if ok: self._run_llm(self.config['cloud_ai']['base_url'], self.config['cloud_ai']['api_key'], m, "cloud_ai")
def _run_llm(self, url, key, model, provider="default"):
p = self.prompt_edit.toPlainText().strip()
if not p: return
self.prompt_edit.setReadOnly(True)
self.llm_worker = LLMWorker(url, key, model, p, provider)
self.llm_worker.finished_signal.connect(self._on_llm_ok)
self.llm_worker.error_signal.connect(self._on_llm_err)
self.llm_worker.start()
def _on_llm_ok(self, t):
self.prompt_edit.setPlainText(t); self.prompt_edit.setReadOnly(False)
def _on_llm_err(self, e):
QMessageBox.warning(self, "AI Error", e); self.prompt_edit.setReadOnly(False)
if __name__ == "__main__":
app = QApplication(sys.argv)
app.setFont(QFont("Segoe UI", 10))
window = MainWindow(); window.show()
sys.exit(app.exec_())