AI 아트 갤러리 구축[미니프로젝트]

Taixi·2024년 12월 30일

9oormthon DEEP DIVE IN Goorm Stable Diffusion 미니팀프로젝트

생성형 AI 교육

목록 보기

34/35

필요한 패키지 설치

!pip install flask torch diffusers transformers Pillow pyngrok

필요한 라이브러리 임포트

from flask import Flask, request, render_template, redirect, url_for
from pyngrok import ngrok, conf
import torch
from diffusers import StableDiffusionPipeline
import os
from threading import Thread
import logging

로깅 설정

logging.basicConfig(level=logging.DEBUG)

Flask 애플리케이션 코드

app = Flask(__name__)

에러 핸들러 추가

@app.errorhandler(Exception)
def handle_error(error):
    app.logger.error(f'An error occurred: {error}')
    return str(error), 500

Stable Diffusion 모델 로드

try:
    model_id = "CompVis/stable-diffusion-v1-4"
    device = "cuda" if torch.cuda.is_available() else "cpu"
    pipe = StableDiffusionPipeline.from_pretrained(
        model_id,
        torch_dtype=torch.float32,  
        safety_checker=None  
    ).to(device)
except Exception as e:
    app.logger.error(f'Model loading error: {e}')
    raise

이미지 생성 함수

def generate_image(prompt):
    try:
        image = pipe(prompt).images[0]
        filename = f"{prompt.replace(' ', '_')}_{os.urandom(4).hex()}.png"
        image_path = os.path.join('static/images', filename)
        image.save(image_path)
        return filename
    except Exception as e:
        app.logger.error(f'Image generation error: {e}')
        raise

라우트 설정

@app.route('/', methods=['GET', 'POST'])
def index():
    try:
        if request.method == 'POST':
            prompt = request.form.get('prompt')
            if not prompt:
                return "Prompt is required", 400

            image_path = generate_image(prompt)
            return redirect(url_for('gallery', image_path=image_path))
        return render_template('index.html')
    except Exception as e:
        app.logger.error(f'Index route error: {e}')
        raise

@app.route('/gallery')
def gallery():
    try:
        image_path = request.args.get('image_path')
        images = os.listdir('static/images')
        return render_template('gallery.html', images=images, new_image=image_path)
    except Exception as e:
        app.logger.error(f'Gallery route error: {e}')
        raise

디렉토리 생성

os.makedirs('static/images', exist_ok=True)
os.makedirs('templates', exist_ok=True)

HTML 템플릿 생성

with open('templates/index.html', 'w') as f:
    f.write('''
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <title>Online Art Gallery</title>
        <style>
            body { font-family: Arial; max-width: 800px; margin: 0 auto; padding: 20px; }
            form { margin: 20px 0; }
            input[type="text"] { width: 100%; padding: 10px; margin: 10px 0; }
            button { padding: 10px 20px; background: #007bff; color: white; border: none; }
        </style>
    </head>
    <body>
        <h1>Stable Diffusion Art Gallery</h1>
        <form method="post">
            <label for="prompt">Enter a prompt to generate an image:</label>
            <input type="text" id="prompt" name="prompt" required>
            <button type="submit">Generate</button>
        </form>
    </body>
    </html>
    ''')

with open('templates/gallery.html', 'w') as f:
    f.write('''
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <title>Gallery</title>
        <style>
            body { font-family: Arial; max-width: 800px; margin: 0 auto; padding: 20px; }
            img { max-width: 100%; margin: 10px 0; }
            .gallery { display: grid; grid-template-columns: repeat(auto-fill, minmax(200px, 1fr)); gap: 20px; }
        </style>
    </head>
    <body>
        <h1>Generated Images</h1>
        {% if new_image %}
            <h2>New Image</h2>
            <img src="{{ url_for('static', filename='images/' + new_image) }}" alt="New Image">
        {% endif %}
        <h2>All Images</h2>
        <div class="gallery">
            {% for image in images %}
                <img src="{{ url_for('static', filename='images/' + image) }}" alt="Image">
            {% endfor %}
        </div>
        <p><a href="/">Generate New Image</a></p>
    </body>
    </html>
    ''')

ngrok 설정 및 실행

def run_app():
    app.run(debug=True, port=5000)

ngrok.set_auth_token("사용자 아이디")
public_url = ngrok.connect(5000)
print(f'Public URL: {public_url.public_url}')

thread = Thread(target=run_app)
thread.daemon = True
thread.start()