app.py

import os
import random
import requests
import json
from flask import Flask, render_template, request, redirect, url_for, flash, send_from_directory, session
from dotenv import load_dotenv

load_dotenv()  # Load environment variables from .env file

# Load auth credentials
AUTH_CREDENTIALS = json.loads(os.getenv('PATPIX_USER', '{"admin":"password"}'))

app = Flask(__name__)
app.secret_key = os.getenv('SECRET_KEY', 'your_secret_key_here')

# Configuration
UPLOAD_FOLDER = 'outputs'
ALLOWED_EXTENSIONS = {'jpg'}
MODELS = ["flux-1-dev-fp8", "flux-1-schnell-fp8"]

# Ensure output directory exists
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def generate_image(api_key, model, prompt):
    url = f"https://api.fireworks.ai/inference/v1/workflows/accounts/fireworks/models/{model}/text_to_image"
    headers = {
        "Content-Type": "application/json",
        "Accept": "image/jpeg",
        "Authorization": f"Bearer {api_key}",
    }
    data = {
        "prompt": prompt,
        "aspect_ratio": "16:9",
        "guidance_scale": 3.5,
        "num_inference_steps": 30,
        "seed": random.randint(0,1000000)
    }
    response = requests.post(url, headers=headers, json=data)
    if response.status_code == 200:
        return response.content
    else:
        raise Exception(f"API Error: {response.status_code} - {response.text}")

def login_required(f):
    def decorated_function(*args, **kwargs):
        if 'username' not in session or session['username'] not in AUTH_CREDENTIALS:
            return redirect(url_for('login'))
        return f(*args, **kwargs)
    decorated_function.__name__ = f.__name__
    return decorated_function

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        
        if username in AUTH_CREDENTIALS and AUTH_CREDENTIALS[username] == password:
            session['username'] = username
            return redirect(url_for('index'))
        else:
            flash('Invalid username or password', 'error')
    return render_template('login.html')

@app.route('/logout')
def logout():
    session.pop('username', None)
    return redirect(url_for('login'))

@app.route('/', methods=['GET', 'POST'])
@login_required
def index():
    if request.method == 'POST':
        api_key = os.getenv('FIREWORKS_API_KEY')
        if not api_key:
            flash('API key not configured in server', 'error')
            return redirect(url_for('index'))
        model = request.form.get('model')
        prompt = request.form.get('prompt')
        count = int(request.form.get('count', 1))

        if not api_key:
            flash('Please enter your API key', 'error')
            return redirect(url_for('index'))

        if not prompt:
            flash('Please enter a prompt', 'error')
            return redirect(url_for('index'))

        try:
            generated_files = []
            for i in range(count):
                image_data = generate_image(api_key, model, prompt)
                if image_data:
                    filename = f"image_{i+1}.jpg"
                    filepath = os.path.join(UPLOAD_FOLDER, filename)
                    with open(filepath, 'wb') as f:
                        f.write(image_data)
                    generated_files.append(filename)
            
            flash(f'Successfully generated {len(generated_files)} images!', 'success')
            return render_template('index.html', models=MODELS, generated_files=generated_files, prompt=prompt)
        
        except Exception as e:
            flash(f'Error generating images: {str(e)}', 'error')
            return redirect(url_for('index'))

    return render_template('index.html', models=MODELS)

@app.route('/outputs/<filename>')
@login_required
def download_file(filename):
    return send_from_directory(UPLOAD_FOLDER, filename)

if __name__ == '__main__':
    app.run(debug=True)