import json import os import re import logging from flask import render_template, request, redirect, url_for, flash, session, current_app from models import db, Character, Look, Outfit, Action, Style, Scene, Detailer, Settings from sqlalchemy.orm.attributes import flag_modified from utils import _LORA_DEFAULTS from services.llm import load_prompt, call_llm logger = logging.getLogger('gaze') _RESOURCE_TRANSFER_MAP = { 'looks': { 'model_class': Look, 'id_field': 'look_id', 'target_dir': 'LOOKS_DIR', 'index_route': 'looks_index', 'detail_route': 'look_detail', 'name_key': 'look_name', 'id_key': 'look_id', }, 'outfits': { 'model_class': Outfit, 'id_field': 'outfit_id', 'target_dir': 'CLOTHING_DIR', 'index_route': 'outfits_index', 'detail_route': 'outfit_detail', 'name_key': 'outfit_name', 'id_key': 'outfit_id', }, 'actions': { 'model_class': Action, 'id_field': 'action_id', 'target_dir': 'ACTIONS_DIR', 'index_route': 'actions_index', 'detail_route': 'action_detail', 'name_key': 'action_name', 'id_key': 'action_id', }, 'styles': { 'model_class': Style, 'id_field': 'style_id', 'target_dir': 'STYLES_DIR', 'index_route': 'styles_index', 'detail_route': 'style_detail', 'name_key': 'style_name', 'id_key': 'style_id', }, 'scenes': { 'model_class': Scene, 'id_field': 'scene_id', 'target_dir': 'SCENES_DIR', 'index_route': 'scenes_index', 'detail_route': 'scene_detail', 'name_key': 'scene_name', 'id_key': 'scene_id', }, 'detailers': { 'model_class': Detailer, 'id_field': 'detailer_id', 'target_dir': 'DETAILERS_DIR', 'index_route': 'detailers_index', 'detail_route': 'detailer_detail', 'name_key': 'detailer_name', 'id_key': 'detailer_id', }, } _TRANSFER_TARGET_CATEGORIES = ['looks', 'outfits', 'actions', 'styles', 'scenes', 'detailers'] def register_routes(app): def _create_minimal_template(target_category, slug, name, source_data, source_name='resource'): """Create a minimal template for the target category with basic fields.""" templates = { 'looks': { 'look_id': slug, 'look_name': name, 'positive': source_data.get('positive', ''), 'negative': source_data.get('negative', ''), 'lora': source_data.get('lora', {'lora_name': '', 'lora_weight': 0.8, 'lora_triggers': ''}), 'tags': source_data.get('tags', []), 'description': f"Transferred from {source_name}" }, 'outfits': { 'outfit_id': slug, 'outfit_name': name, 'wardrobe': source_data.get('wardrobe', { 'base': '', 'head': '', 'upper_body': '', 'lower_body': '', 'hands': '', 'feet': '', 'additional': '' }), 'lora': source_data.get('lora', {'lora_name': '', 'lora_weight': 0.8, 'lora_triggers': ''}), 'tags': source_data.get('tags', []), 'description': f"Transferred from {source_name}" }, 'actions': { 'action_id': slug, 'action_name': name, 'action': source_data.get('action', { 'base': '', 'head': '', 'upper_body': '', 'lower_body': '', 'hands': '', 'feet': '', 'additional': '' }), 'lora': source_data.get('lora', {'lora_name': '', 'lora_weight': 1.0, 'lora_triggers': ''}), 'tags': source_data.get('tags', []), 'description': f"Transferred from {source_name}" }, 'styles': { 'style_id': slug, 'style_name': name, 'style': source_data.get('style', {'artist_name': '', 'artistic_style': ''}), 'lora': source_data.get('lora', {'lora_name': '', 'lora_weight': 1.0, 'lora_triggers': ''}), 'tags': source_data.get('tags', []), 'description': f"Transferred from {source_name}" }, 'scenes': { 'scene_id': slug, 'scene_name': name, 'scene': source_data.get('scene', { 'background': '', 'foreground': '', 'furniture': [], 'colors': [], 'lighting': '', 'theme': '' }), 'lora': source_data.get('lora', {'lora_name': '', 'lora_weight': 1.0, 'lora_triggers': ''}), 'tags': source_data.get('tags', []), 'description': f"Transferred from {source_name}" }, 'detailers': { 'detailer_id': slug, 'detailer_name': name, 'prompt': source_data.get('prompt', source_data.get('positive', '')), 'lora': source_data.get('lora', {'lora_name': '', 'lora_weight': 1.0, 'lora_triggers': ''}), 'tags': source_data.get('tags', []), 'description': f"Transferred from {source_name}" }, } return templates.get(target_category, { f'{target_category.rstrip("s")}_id': slug, f'{target_category.rstrip("s")}_name': name, 'description': f"Transferred from {source_name}", 'tags': source_data.get('tags', []), 'lora': source_data.get('lora', {}) }) @app.route('/resource///transfer', methods=['GET', 'POST']) def transfer_resource(category, slug): """Generic resource transfer route.""" if category not in _RESOURCE_TRANSFER_MAP: flash(f'Invalid category: {category}') return redirect(url_for('index')) source_config = _RESOURCE_TRANSFER_MAP[category] model_class = source_config['model_class'] id_field = source_config['id_field'] resource = model_class.query.filter_by(slug=slug).first_or_404() resource_name = resource.name resource_data = resource.data if request.method == 'POST': target_category = request.form.get('target_category') new_name = request.form.get('new_name', '').strip() new_id = request.form.get('new_id', '').strip() use_llm = request.form.get('use_llm') in ('on', 'true', '1', 'yes') or request.form.get('use_llm') is not None transfer_lora = request.form.get('transfer_lora') == 'on' remove_original = request.form.get('remove_original') == 'on' if not new_name: flash('New name is required for transfer') return redirect(url_for('transfer_resource', category=category, slug=slug)) if len(new_name) > 100: flash('New name must be 100 characters or less') return redirect(url_for('transfer_resource', category=category, slug=slug)) if not new_id: new_id = re.sub(r'[^a-zA-Z0-9]+', '_', new_name.lower()).strip('_') safe_slug = re.sub(r'[^a-zA-Z0-9_]', '', new_id) if not safe_slug: safe_slug = 'transferred' if target_category not in _RESOURCE_TRANSFER_MAP: flash('Invalid target category') return redirect(url_for('transfer_resource', category=category, slug=slug)) if target_category == category: flash('Cannot transfer to the same category') return redirect(url_for('transfer_resource', category=category, slug=slug)) target_config = _RESOURCE_TRANSFER_MAP[target_category] target_model_class = target_config['model_class'] target_id_field = target_config['id_field'] target_dir = current_app.config[target_config['target_dir']] target_name_key = target_config['name_key'] target_id_key = target_config['id_key'] base_slug = safe_slug counter = 1 while os.path.exists(os.path.join(target_dir, f"{safe_slug}.json")): safe_slug = f"{base_slug}_{counter}" counter += 1 if use_llm: try: category_system_prompts = { 'outfits': 'outfit_system.txt', 'actions': 'action_system.txt', 'styles': 'style_system.txt', 'scenes': 'scene_system.txt', 'detailers': 'detailer_system.txt', 'looks': 'look_system.txt', } system_prompt_file = category_system_prompts.get(target_category) system_prompt = load_prompt(system_prompt_file) if system_prompt_file else None if not system_prompt: system_prompt = load_prompt('transfer_system.txt') if not system_prompt: system_prompt = f"""You are an AI assistant that creates {target_category.rstrip('s')} profiles for AI image generation. Your task is to create a {target_category.rstrip('s')} profile based on the source resource data provided. Target type: {target_category} Required JSON Structure: - {target_id_key}: "{safe_slug}" - {target_name_key}: "{new_name}" """ source_summary = json.dumps(resource_data, indent=2) llm_prompt = f"""Create a {target_category.rstrip('s')} profile named "{new_name}" (ID: {safe_slug}) based on this source {category.rstrip('s')} data: Source {category.rstrip('s')} name: {resource_name} Source data: {source_summary} Generate a complete {target_category.rstrip('s')} profile with all required fields for the {target_category} category.""" llm_response = call_llm(llm_prompt, system_prompt) clean_json = llm_response.replace('```json', '').replace('```', '').strip() new_data = json.loads(clean_json) new_data[target_id_key] = safe_slug new_data[target_name_key] = new_name except Exception as e: logger.exception(f"LLM transfer error: {e}") flash(f'Failed to generate {target_category.rstrip("s")} with AI: {e}') return redirect(url_for('transfer_resource', category=category, slug=slug)) else: new_data = _create_minimal_template(target_category, safe_slug, new_name, resource_data, resource_name) try: settings = Settings.query.first() lora_moved = False if transfer_lora and 'lora' in resource_data and resource_data['lora'].get('lora_name'): old_lora_path = resource_data['lora']['lora_name'] lora_dir_map = { 'characters': getattr(settings, 'lora_dir_characters', None) or _LORA_DEFAULTS.get('characters', ''), 'looks': getattr(settings, 'lora_dir_looks', None) or getattr(settings, 'lora_dir_characters', None) or _LORA_DEFAULTS.get('looks', _LORA_DEFAULTS.get('characters', '')), 'outfits': getattr(settings, 'lora_dir_outfits', None) or _LORA_DEFAULTS.get('outfits', ''), 'actions': getattr(settings, 'lora_dir_actions', None) or _LORA_DEFAULTS.get('actions', ''), 'styles': getattr(settings, 'lora_dir_styles', None) or _LORA_DEFAULTS.get('styles', ''), 'scenes': getattr(settings, 'lora_dir_scenes', None) or _LORA_DEFAULTS.get('scenes', ''), 'detailers': getattr(settings, 'lora_dir_detailers', None) or _LORA_DEFAULTS.get('detailers', ''), } target_lora_dir = lora_dir_map.get(target_category) source_lora_dir = lora_dir_map.get(category) if old_lora_path and target_lora_dir and source_lora_dir: lora_filename = os.path.basename(old_lora_path) abs_source_path = os.path.join('/ImageModels/lora', old_lora_path) if not os.path.exists(abs_source_path): abs_source_path = os.path.join(source_lora_dir, lora_filename) abs_target_path = os.path.join(target_lora_dir, lora_filename) if os.path.exists(abs_source_path): try: import shutil os.makedirs(target_lora_dir, exist_ok=True) shutil.move(abs_source_path, abs_target_path) target_subfolder = os.path.basename(target_lora_dir.rstrip('/')) new_data['lora']['lora_name'] = f"Illustrious/{target_subfolder}/{lora_filename}" lora_moved = True flash(f'Moved LoRA file to {target_lora_dir}') except Exception as lora_e: logger.exception(f"LoRA move error: {lora_e}") flash(f'Warning: Failed to move LoRA file: {lora_e}', 'warning') else: flash(f'Warning: Source LoRA file not found at {abs_source_path}', 'warning') file_path = os.path.join(target_dir, f"{safe_slug}.json") with open(file_path, 'w') as f: json.dump(new_data, f, indent=2) new_entity = target_model_class( **{target_id_field: safe_slug}, slug=safe_slug, filename=f"{safe_slug}.json", name=new_name, data=new_data ) db.session.add(new_entity) if remove_original: try: source_dir = current_app.config[source_config['target_dir']] orig_file_path = os.path.join(source_dir, resource.filename or f"{resource.slug}.json") if os.path.exists(orig_file_path): os.remove(orig_file_path) db.session.delete(resource) flash(f'Removed original {category.rstrip("s")}: {resource_name}') except Exception as rm_e: logger.exception(f"Error removing original: {rm_e}") flash(f'Warning: Failed to remove original: {rm_e}', 'warning') db.session.commit() flash(f'Successfully transferred to {target_category.rstrip("s")}: {new_name}') return redirect(url_for(target_config['index_route'], highlight=safe_slug)) except Exception as e: logger.exception(f"Transfer save error: {e}") flash(f'Failed to save transferred {target_category.rstrip("s")}: {e}') return redirect(url_for('transfer_resource', category=category, slug=slug)) available_targets = [(cat, cat.rstrip('s').replace('_', ' ').title()) for cat in _TRANSFER_TARGET_CATEGORIES if cat != category] return render_template('transfer_resource.html', category=category, resource=resource, available_targets=available_targets, cancel_route=source_config['detail_route'])