598 lines
24 KiB
Python
598 lines
24 KiB
Python
import os
|
|
import base64
|
|
import uuid
|
|
import glob
|
|
import json
|
|
from datetime import datetime
|
|
from io import BytesIO
|
|
from send2trash import send2trash
|
|
from flask import Flask, render_template, request, jsonify, url_for
|
|
from google import genai
|
|
from google.genai import types
|
|
from PIL import Image, PngImagePlugin
|
|
|
|
app = Flask(__name__)
|
|
app.config['SEND_FILE_MAX_AGE_DEFAULT'] = 0
|
|
|
|
# Ensure generated directory exists inside Flask static folder
|
|
GENERATED_DIR = os.path.join(app.static_folder, 'generated')
|
|
os.makedirs(GENERATED_DIR, exist_ok=True)
|
|
|
|
# Ensure uploads directory exists
|
|
UPLOADS_DIR = os.path.join(app.static_folder, 'uploads')
|
|
os.makedirs(UPLOADS_DIR, exist_ok=True)
|
|
|
|
@app.route('/')
|
|
def index():
|
|
return render_template('index.html')
|
|
|
|
@app.route('/generate', methods=['POST'])
|
|
def generate_image():
|
|
multipart = request.content_type and 'multipart/form-data' in request.content_type
|
|
|
|
if multipart:
|
|
form = request.form
|
|
prompt = form.get('prompt')
|
|
aspect_ratio = form.get('aspect_ratio')
|
|
resolution = form.get('resolution', '2K')
|
|
api_key = form.get('api_key') or os.environ.get('GOOGLE_API_KEY')
|
|
reference_files = request.files.getlist('reference_images')
|
|
reference_paths_json = form.get('reference_image_paths')
|
|
else:
|
|
data = request.get_json() or {}
|
|
prompt = data.get('prompt')
|
|
aspect_ratio = data.get('aspect_ratio')
|
|
resolution = data.get('resolution', '2K')
|
|
api_key = data.get('api_key') or os.environ.get('GOOGLE_API_KEY')
|
|
reference_files = []
|
|
reference_paths_json = data.get('reference_image_paths')
|
|
|
|
if not prompt:
|
|
return jsonify({'error': 'Prompt is required'}), 400
|
|
|
|
if not api_key:
|
|
return jsonify({'error': 'API Key is required.'}), 401
|
|
|
|
try:
|
|
client = genai.Client(api_key=api_key)
|
|
|
|
image_config_args = {
|
|
"image_size": resolution
|
|
}
|
|
|
|
if aspect_ratio and aspect_ratio != 'Auto':
|
|
image_config_args["aspect_ratio"] = aspect_ratio
|
|
|
|
# Process reference paths and files
|
|
final_reference_paths = []
|
|
contents = [prompt]
|
|
|
|
# Parse reference paths from frontend
|
|
frontend_paths = []
|
|
if reference_paths_json:
|
|
try:
|
|
frontend_paths = json.loads(reference_paths_json)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
# If no paths provided but we have files (legacy or simple upload), treat all as new uploads
|
|
# But we need to handle the mix.
|
|
# Strategy: Iterate frontend_paths. If it looks like a path/URL, keep it.
|
|
# If it doesn't (or is null), consume from reference_files.
|
|
|
|
file_index = 0
|
|
|
|
# If frontend_paths is empty but we have files, just use the files
|
|
if not frontend_paths and reference_files:
|
|
for _ in reference_files:
|
|
frontend_paths.append(None) # Placeholder for each file
|
|
|
|
for path in frontend_paths:
|
|
if path and (path.startswith('/') or path.startswith('http')):
|
|
# Existing path/URL
|
|
final_reference_paths.append(path)
|
|
# We also need to add the image content to the prompt
|
|
# We need to fetch it or read it if it's local (server-side local)
|
|
# If it's a URL we generated, it's in static/generated or static/uploads
|
|
# path might be "http://localhost:8888/static/generated/..." or "/static/generated/..."
|
|
|
|
# Extract relative path to open file
|
|
# Assuming path contains '/static/'
|
|
try:
|
|
if '/static/' in path:
|
|
rel_path = path.split('/static/')[1]
|
|
abs_path = os.path.join(app.static_folder, rel_path)
|
|
if os.path.exists(abs_path):
|
|
img = Image.open(abs_path)
|
|
contents.append(img)
|
|
else:
|
|
print(f"Warning: Reference file not found at {abs_path}")
|
|
else:
|
|
print(f"Warning: Could not resolve local path for {path}")
|
|
except Exception as e:
|
|
print(f"Error loading reference from path {path}: {e}")
|
|
|
|
elif file_index < len(reference_files):
|
|
# New upload
|
|
file = reference_files[file_index]
|
|
file_index += 1
|
|
|
|
try:
|
|
# Save to uploads
|
|
ext = os.path.splitext(file.filename)[1]
|
|
if not ext:
|
|
ext = '.png'
|
|
filename = f"{uuid.uuid4()}{ext}"
|
|
filepath = os.path.join(UPLOADS_DIR, filename)
|
|
|
|
# We need to read the file for Gemini AND save it
|
|
# file.stream is a stream.
|
|
file.stream.seek(0)
|
|
file_bytes = file.read()
|
|
|
|
with open(filepath, 'wb') as f:
|
|
f.write(file_bytes)
|
|
|
|
# Add to contents
|
|
image = Image.open(BytesIO(file_bytes))
|
|
contents.append(image)
|
|
|
|
# Add to final paths
|
|
# URL for the uploaded file
|
|
rel_path = os.path.join('uploads', filename)
|
|
file_url = url_for('static', filename=rel_path)
|
|
final_reference_paths.append(file_url)
|
|
|
|
except Exception as e:
|
|
print(f"Error processing uploaded file: {e}")
|
|
continue
|
|
|
|
model_name = "gemini-3-pro-image-preview"
|
|
response = client.models.generate_content(
|
|
model=model_name,
|
|
contents=contents,
|
|
config=types.GenerateContentConfig(
|
|
response_modalities=['IMAGE'],
|
|
image_config=types.ImageConfig(**image_config_args),
|
|
)
|
|
)
|
|
|
|
for part in response.parts:
|
|
if part.inline_data:
|
|
image_bytes = part.inline_data.data
|
|
|
|
image = Image.open(BytesIO(image_bytes))
|
|
png_info = PngImagePlugin.PngInfo()
|
|
|
|
date_str = datetime.now().strftime("%Y%m%d")
|
|
|
|
# Find existing files to determine next ID
|
|
search_pattern = os.path.join(GENERATED_DIR, f"{model_name}_{date_str}_*.png")
|
|
existing_files = glob.glob(search_pattern)
|
|
max_id = 0
|
|
for f in existing_files:
|
|
try:
|
|
basename = os.path.basename(f)
|
|
name_without_ext = os.path.splitext(basename)[0]
|
|
id_part = name_without_ext.split('_')[-1]
|
|
id_num = int(id_part)
|
|
if id_num > max_id:
|
|
max_id = id_num
|
|
except ValueError:
|
|
continue
|
|
|
|
next_id = max_id + 1
|
|
filename = f"{model_name}_{date_str}_{next_id}.png"
|
|
filepath = os.path.join(GENERATED_DIR, filename)
|
|
rel_path = os.path.join('generated', filename)
|
|
image_url = url_for('static', filename=rel_path)
|
|
|
|
metadata = {
|
|
'prompt': prompt,
|
|
'aspect_ratio': aspect_ratio or 'Auto',
|
|
'resolution': resolution,
|
|
'reference_images': final_reference_paths,
|
|
}
|
|
|
|
png_info.add_text('sdvn_meta', json.dumps(metadata))
|
|
|
|
buffer = BytesIO()
|
|
image.save(buffer, format='PNG', pnginfo=png_info)
|
|
final_bytes = buffer.getvalue()
|
|
|
|
# Save image to file
|
|
with open(filepath, 'wb') as f:
|
|
f.write(final_bytes)
|
|
|
|
image_data = base64.b64encode(final_bytes).decode('utf-8')
|
|
return jsonify({
|
|
'image': image_url,
|
|
'image_data': image_data,
|
|
'metadata': metadata,
|
|
})
|
|
|
|
return jsonify({'error': 'No image generated'}), 500
|
|
|
|
except Exception as e:
|
|
return jsonify({'error': str(e)}), 500
|
|
|
|
@app.route('/delete_image', methods=['POST'])
|
|
def delete_image():
|
|
data = request.get_json()
|
|
filename = data.get('filename')
|
|
|
|
if not filename:
|
|
return jsonify({'error': 'Filename is required'}), 400
|
|
|
|
# Security check: ensure filename is just a basename, no paths
|
|
filename = os.path.basename(filename)
|
|
filepath = os.path.join(GENERATED_DIR, filename)
|
|
|
|
if os.path.exists(filepath):
|
|
try:
|
|
send2trash(filepath)
|
|
return jsonify({'success': True})
|
|
except Exception as e:
|
|
return jsonify({'error': str(e)}), 500
|
|
else:
|
|
return jsonify({'error': 'File not found'}), 404
|
|
|
|
@app.route('/gallery')
|
|
def get_gallery():
|
|
# List all png files in generated dir, sorted by modification time (newest first)
|
|
files = glob.glob(os.path.join(GENERATED_DIR, '*.png'))
|
|
files.sort(key=os.path.getmtime, reverse=True)
|
|
|
|
image_urls = [url_for('static', filename=f'generated/{os.path.basename(f)}') for f in files]
|
|
response = jsonify({'images': image_urls})
|
|
response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate"
|
|
return response
|
|
|
|
@app.route('/prompts')
|
|
def get_prompts():
|
|
category = request.args.get('category')
|
|
|
|
try:
|
|
all_prompts = []
|
|
|
|
# Read prompts.json file
|
|
prompts_path = os.path.join(os.path.dirname(__file__), 'prompts.json')
|
|
if os.path.exists(prompts_path):
|
|
with open(prompts_path, 'r', encoding='utf-8') as f:
|
|
all_prompts.extend(json.load(f))
|
|
|
|
# Read user_prompts.json file and mark as user templates
|
|
user_prompts_path = os.path.join(os.path.dirname(__file__), 'user_prompts.json')
|
|
if os.path.exists(user_prompts_path):
|
|
try:
|
|
with open(user_prompts_path, 'r', encoding='utf-8') as f:
|
|
user_prompts = json.load(f)
|
|
if isinstance(user_prompts, list):
|
|
# Mark each user template and add index for editing
|
|
for idx, template in enumerate(user_prompts):
|
|
template['isUserTemplate'] = True
|
|
template['userTemplateIndex'] = idx
|
|
all_prompts.extend(user_prompts)
|
|
except json.JSONDecodeError:
|
|
pass # Ignore if empty or invalid
|
|
|
|
# Filter by category if specified
|
|
if category:
|
|
all_prompts = [p for p in all_prompts if p.get('category') == category]
|
|
|
|
response = jsonify({'prompts': all_prompts})
|
|
response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate"
|
|
return response
|
|
except Exception as e:
|
|
return jsonify({'error': str(e)}), 500
|
|
|
|
@app.route('/save_template', methods=['POST'])
|
|
def save_template():
|
|
try:
|
|
import requests
|
|
from urllib.parse import urlparse
|
|
|
|
# Handle multipart form data
|
|
title = request.form.get('title')
|
|
prompt = request.form.get('prompt')
|
|
mode = request.form.get('mode', 'generate')
|
|
category = request.form.get('category', 'User')
|
|
|
|
if not title or not prompt:
|
|
return jsonify({'error': 'Title and prompt are required'}), 400
|
|
|
|
# Handle preview image
|
|
preview_path = None
|
|
preview_dir = os.path.join(app.static_folder, 'preview')
|
|
os.makedirs(preview_dir, exist_ok=True)
|
|
|
|
# Check if file was uploaded
|
|
if 'preview' in request.files:
|
|
file = request.files['preview']
|
|
if file.filename:
|
|
ext = os.path.splitext(file.filename)[1]
|
|
if not ext:
|
|
ext = '.png'
|
|
|
|
filename = f"template_{uuid.uuid4()}{ext}"
|
|
filepath = os.path.join(preview_dir, filename)
|
|
file.save(filepath)
|
|
|
|
preview_path = url_for('static', filename=f'preview/{filename}')
|
|
|
|
# If no file uploaded, check if URL/path provided
|
|
if not preview_path:
|
|
preview_url = request.form.get('preview_path')
|
|
if preview_url:
|
|
try:
|
|
# Check if it's a URL or local path
|
|
if preview_url.startswith('http://') or preview_url.startswith('https://'):
|
|
# Download from URL
|
|
response = requests.get(preview_url, timeout=10)
|
|
response.raise_for_status()
|
|
|
|
# Determine extension from content-type or URL
|
|
content_type = response.headers.get('content-type', '')
|
|
if 'image/png' in content_type:
|
|
ext = '.png'
|
|
elif 'image/jpeg' in content_type or 'image/jpg' in content_type:
|
|
ext = '.jpg'
|
|
elif 'image/webp' in content_type:
|
|
ext = '.webp'
|
|
else:
|
|
# Try to get from URL
|
|
parsed = urlparse(preview_url)
|
|
ext = os.path.splitext(parsed.path)[1] or '.png'
|
|
|
|
filename = f"template_{uuid.uuid4()}{ext}"
|
|
filepath = os.path.join(preview_dir, filename)
|
|
|
|
with open(filepath, 'wb') as f:
|
|
f.write(response.content)
|
|
|
|
preview_path = url_for('static', filename=f'preview/{filename}')
|
|
|
|
elif preview_url.startswith('/static/'):
|
|
# Local path - copy to preview folder
|
|
rel_path = preview_url.split('/static/')[1]
|
|
source_path = os.path.join(app.static_folder, rel_path)
|
|
|
|
if os.path.exists(source_path):
|
|
ext = os.path.splitext(source_path)[1] or '.png'
|
|
filename = f"template_{uuid.uuid4()}{ext}"
|
|
dest_path = os.path.join(preview_dir, filename)
|
|
|
|
import shutil
|
|
shutil.copy2(source_path, dest_path)
|
|
|
|
preview_path = url_for('static', filename=f'preview/{filename}')
|
|
else:
|
|
# File doesn't exist, use original path
|
|
preview_path = preview_url
|
|
else:
|
|
# Use as-is if it's already a valid path
|
|
preview_path = preview_url
|
|
|
|
except Exception as e:
|
|
print(f"Error processing preview image URL: {e}")
|
|
# Use the original URL if processing fails
|
|
preview_path = preview_url
|
|
|
|
new_template = {
|
|
'title': title,
|
|
'prompt': prompt,
|
|
'mode': mode,
|
|
'category': category,
|
|
'preview': preview_path
|
|
}
|
|
|
|
# Save to user_prompts.json
|
|
user_prompts_path = os.path.join(os.path.dirname(__file__), 'user_prompts.json')
|
|
user_prompts = []
|
|
|
|
if os.path.exists(user_prompts_path):
|
|
try:
|
|
with open(user_prompts_path, 'r', encoding='utf-8') as f:
|
|
content = f.read()
|
|
if content.strip():
|
|
user_prompts = json.loads(content)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
user_prompts.append(new_template)
|
|
|
|
with open(user_prompts_path, 'w', encoding='utf-8') as f:
|
|
json.dump(user_prompts, f, indent=4, ensure_ascii=False)
|
|
|
|
return jsonify({'success': True, 'template': new_template})
|
|
|
|
except Exception as e:
|
|
print(f"Error saving template: {e}")
|
|
return jsonify({'error': str(e)}), 500
|
|
|
|
@app.route('/update_template', methods=['POST'])
|
|
def update_template():
|
|
try:
|
|
import requests
|
|
from urllib.parse import urlparse
|
|
|
|
# Get template index
|
|
template_index = request.form.get('template_index')
|
|
if template_index is None:
|
|
return jsonify({'error': 'Template index is required'}), 400
|
|
|
|
template_index = int(template_index)
|
|
|
|
# Handle multipart form data
|
|
title = request.form.get('title')
|
|
prompt = request.form.get('prompt')
|
|
mode = request.form.get('mode', 'generate')
|
|
category = request.form.get('category', 'User')
|
|
|
|
if not title or not prompt:
|
|
return jsonify({'error': 'Title and prompt are required'}), 400
|
|
|
|
# Handle preview image (same logic as save_template)
|
|
preview_path = None
|
|
preview_dir = os.path.join(app.static_folder, 'preview')
|
|
os.makedirs(preview_dir, exist_ok=True)
|
|
|
|
# Check if file was uploaded
|
|
if 'preview' in request.files:
|
|
file = request.files['preview']
|
|
if file.filename:
|
|
ext = os.path.splitext(file.filename)[1]
|
|
if not ext:
|
|
ext = '.png'
|
|
|
|
filename = f"template_{uuid.uuid4()}{ext}"
|
|
filepath = os.path.join(preview_dir, filename)
|
|
file.save(filepath)
|
|
|
|
preview_path = url_for('static', filename=f'preview/{filename}')
|
|
|
|
# If no file uploaded, check if URL/path provided
|
|
if not preview_path:
|
|
preview_url = request.form.get('preview_path')
|
|
if preview_url:
|
|
try:
|
|
if preview_url.startswith('http://') or preview_url.startswith('https://'):
|
|
response = requests.get(preview_url, timeout=10)
|
|
response.raise_for_status()
|
|
|
|
content_type = response.headers.get('content-type', '')
|
|
if 'image/png' in content_type:
|
|
ext = '.png'
|
|
elif 'image/jpeg' in content_type or 'image/jpg' in content_type:
|
|
ext = '.jpg'
|
|
elif 'image/webp' in content_type:
|
|
ext = '.webp'
|
|
else:
|
|
parsed = urlparse(preview_url)
|
|
ext = os.path.splitext(parsed.path)[1] or '.png'
|
|
|
|
filename = f"template_{uuid.uuid4()}{ext}"
|
|
filepath = os.path.join(preview_dir, filename)
|
|
|
|
with open(filepath, 'wb') as f:
|
|
f.write(response.content)
|
|
|
|
preview_path = url_for('static', filename=f'preview/{filename}')
|
|
|
|
elif preview_url.startswith('/static/'):
|
|
rel_path = preview_url.split('/static/')[1]
|
|
source_path = os.path.join(app.static_folder, rel_path)
|
|
|
|
if os.path.exists(source_path):
|
|
ext = os.path.splitext(source_path)[1] or '.png'
|
|
filename = f"template_{uuid.uuid4()}{ext}"
|
|
dest_path = os.path.join(preview_dir, filename)
|
|
|
|
import shutil
|
|
shutil.copy2(source_path, dest_path)
|
|
|
|
preview_path = url_for('static', filename=f'preview/{filename}')
|
|
else:
|
|
preview_path = preview_url
|
|
else:
|
|
preview_path = preview_url
|
|
|
|
except Exception as e:
|
|
print(f"Error processing preview image URL: {e}")
|
|
preview_path = preview_url
|
|
|
|
# Read existing user templates
|
|
user_prompts_path = os.path.join(os.path.dirname(__file__), 'user_prompts.json')
|
|
user_prompts = []
|
|
|
|
if os.path.exists(user_prompts_path):
|
|
try:
|
|
with open(user_prompts_path, 'r', encoding='utf-8') as f:
|
|
content = f.read()
|
|
if content.strip():
|
|
user_prompts = json.loads(content)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
# Check if index is valid
|
|
if template_index < 0 or template_index >= len(user_prompts):
|
|
return jsonify({'error': 'Invalid template index'}), 400
|
|
|
|
# Get old template to check for old preview image
|
|
old_template = user_prompts[template_index]
|
|
old_preview = old_template.get('preview', '')
|
|
|
|
# Delete old preview image if it exists in the preview folder
|
|
if old_preview and '/preview/' in old_preview:
|
|
try:
|
|
# Extract filename from URL
|
|
old_filename = old_preview.split('/preview/')[-1]
|
|
old_filepath = os.path.join(preview_dir, old_filename)
|
|
|
|
# Delete old file if it exists
|
|
if os.path.exists(old_filepath):
|
|
os.remove(old_filepath)
|
|
print(f"Deleted old preview image: {old_filepath}")
|
|
except Exception as e:
|
|
print(f"Error deleting old preview image: {e}")
|
|
|
|
# Update the template
|
|
user_prompts[template_index] = {
|
|
'title': title,
|
|
'prompt': prompt,
|
|
'mode': mode,
|
|
'category': category,
|
|
'preview': preview_path
|
|
}
|
|
|
|
# Save back to file
|
|
with open(user_prompts_path, 'w', encoding='utf-8') as f:
|
|
json.dump(user_prompts, f, indent=4, ensure_ascii=False)
|
|
|
|
return jsonify({'success': True, 'template': user_prompts[template_index]})
|
|
|
|
except Exception as e:
|
|
print(f"Error updating template: {e}")
|
|
return jsonify({'error': str(e)}), 500
|
|
|
|
@app.route('/refine_prompt', methods=['POST'])
|
|
def refine_prompt():
|
|
data = request.get_json()
|
|
current_prompt = data.get('current_prompt')
|
|
instruction = data.get('instruction')
|
|
api_key = data.get('api_key') or os.environ.get('GOOGLE_API_KEY')
|
|
|
|
if not api_key:
|
|
return jsonify({'error': 'API Key is required.'}), 401
|
|
|
|
if not instruction:
|
|
return jsonify({'error': 'Instruction is required'}), 400
|
|
|
|
try:
|
|
client = genai.Client(api_key=api_key)
|
|
|
|
system_instruction = "You are an expert prompt engineer for image generation AI. Rewrite the prompt to incorporate the user's instruction while maintaining the original intent and improving quality. Return ONLY the new prompt text, no explanations."
|
|
|
|
prompt_content = f"Current prompt: {current_prompt}\nUser instruction: {instruction}\nNew prompt:"
|
|
|
|
print(f"Refining prompt with instruction: {instruction}")
|
|
|
|
response = client.models.generate_content(
|
|
model="gemini-2.5-flash",
|
|
contents=[prompt_content],
|
|
config=types.GenerateContentConfig(
|
|
system_instruction=system_instruction,
|
|
temperature=0.7,
|
|
)
|
|
)
|
|
|
|
if response.text:
|
|
return jsonify({'refined_prompt': response.text.strip()})
|
|
else:
|
|
return jsonify({'error': 'No response from AI'}), 500
|
|
|
|
except Exception as e:
|
|
return jsonify({'error': str(e)}), 500
|
|
|
|
if __name__ == '__main__':
|
|
app.run(debug=True, port=8888)
|