- IT (1442)
Notice
Recent Posts
Recent Comments
반응형
오늘도 공부
Image Generation Service - 완벽 핸즈온 가이드 본문
반응형
🎯 학습 목표
- Stable Diffusion 아키텍처 이해
- Text-to-Image (T2I) 생성 구현
- LoRA/DreamBooth로 커스텀 모델 학습
- ControlNet으로 이미지 제어
- 이미지 편집 (Inpainting, Outpainting)
- 프로덕션급 이미지 생성 서비스 구축
📋 사전 준비
1. 시스템 요구사항
# GPU 필수 (VRAM 8GB 이상 권장)
# CUDA 설치 확인
nvidia-smi
# 또는 CPU로도 가능 (매우 느림)
2. 개발 환경 설정
# 새 프로젝트 디렉토리
mkdir image-generation-service
cd image-generation-service
# 가상환경 생성
python -m venv venv
source venv/bin/activate # Windows: venv\Scripts\activate
# PyTorch 설치 (CUDA 버전에 맞게)
# CUDA 11.8
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
# 또는 CPU 버전
# pip install torch torchvision torchaudio
# Diffusers 및 관련 패키지
pip install diffusers transformers accelerate
pip install safetensors omegaconf
pip install xformers # 선택적, 속도 향상
# 이미지 처리
pip install pillow opencv-python matplotlib
pip install gradio streamlit
# 기타
pip install python-dotenv huggingface_hub
3. Hugging Face 토큰
# https://huggingface.co/settings/tokens 에서 토큰 발급
# 모델 다운로드를 위해 필요
4. 프로젝트 구조
mkdir -p src models utils outputs
mkdir -p data/training data/lora
touch .env app.py requirements.txt
image-generation-service/
├── src/
│ ├── text_to_image.py # 기본 T2I 생성
│ ├── image_to_image.py # I2I 변환
│ ├── inpainting.py # 이미지 인페인팅
│ ├── controlnet_generator.py # ControlNet
│ ├── lora_trainer.py # LoRA 학습
│ └── prompt_enhancer.py # 프롬프트 개선
├── models/ # 다운로드된 모델
├── utils/
│ ├── image_utils.py # 이미지 처리
│ └── model_loader.py # 모델 로딩
├── outputs/ # 생성된 이미지
├── data/
│ ├── training/ # 학습 데이터
│ └── lora/ # LoRA 모델
├── app.py # Streamlit UI
└── .env
.env 파일
HUGGINGFACE_TOKEN=your_hf_token_here
🎨 Step 1: 기본 Text-to-Image 생성 (25분)
utils/model_loader.py
import os
import torch
from diffusers import (
StableDiffusionPipeline,
StableDiffusionXLPipeline,
DPMSolverMultistepScheduler
)
from huggingface_hub import login
from dotenv import load_dotenv
load_dotenv()
class ModelLoader:
"""Stable Diffusion 모델 로더"""
def __init__(self, use_gpu: bool = True):
"""
Args:
use_gpu: GPU 사용 여부
"""
self.device = "cuda" if use_gpu and torch.cuda.is_available() else "cpu"
print(f"🖥️ 디바이스: {self.device}")
if self.device == "cuda":
print(f" GPU: {torch.cuda.get_device_name(0)}")
print(f" VRAM: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
# Hugging Face 로그인
hf_token = os.getenv("HUGGINGFACE_TOKEN")
if hf_token:
login(token=hf_token)
def load_sd15_model(self, model_id: str = "runwayml/stable-diffusion-v1-5"):
"""Stable Diffusion 1.5 모델 로드"""
print(f"📥 모델 로딩 중: {model_id}")
pipe = StableDiffusionPipeline.from_pretrained(
model_id,
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
safety_checker=None # NSFW 체커 비활성화 (선택)
)
# 스케줄러 최적화
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
# GPU로 이동
pipe = pipe.to(self.device)
# 메모리 최적화
if self.device == "cuda":
pipe.enable_attention_slicing() # VRAM 절약
# pipe.enable_xformers_memory_efficient_attention() # xformers 설치 시
print(f"✅ 모델 로드 완료")
return pipe
def load_sdxl_model(self, model_id: str = "stabilityai/stable-diffusion-xl-base-1.0"):
"""Stable Diffusion XL 모델 로드 (고품질, VRAM 많이 필요)"""
print(f"📥 SDXL 모델 로딩 중: {model_id}")
pipe = StableDiffusionXLPipeline.from_pretrained(
model_id,
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
use_safetensors=True
)
pipe = pipe.to(self.device)
if self.device == "cuda":
pipe.enable_attention_slicing()
print(f"✅ SDXL 모델 로드 완료")
return pipe
def load_custom_model(self, model_path: str):
"""로컬 커스텀 모델 로드"""
print(f"📥 커스텀 모델 로딩: {model_path}")
pipe = StableDiffusionPipeline.from_pretrained(
model_path,
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
safety_checker=None
)
pipe = pipe.to(self.device)
print(f"✅ 커스텀 모델 로드 완료")
return pipe
# 테스트
if __name__ == "__main__":
loader = ModelLoader(use_gpu=True)
pipe = loader.load_sd15_model()
print(f"파이프라인 타입: {type(pipe)}")
src/text_to_image.py
import torch
from PIL import Image
import os
from typing import Optional, List
import sys
sys.path.append('..')
from utils.model_loader import ModelLoader
class TextToImageGenerator:
"""Text-to-Image 생성기"""
def __init__(self, model_type: str = "sd15", use_gpu: bool = True):
"""
Args:
model_type: 'sd15' 또는 'sdxl'
use_gpu: GPU 사용 여부
"""
self.loader = ModelLoader(use_gpu=use_gpu)
if model_type == "sd15":
self.pipe = self.loader.load_sd15_model()
elif model_type == "sdxl":
self.pipe = self.loader.load_sdxl_model()
else:
raise ValueError("model_type은 'sd15' 또는 'sdxl'이어야 합니다")
self.device = self.loader.device
def generate(
self,
prompt: str,
negative_prompt: Optional[str] = None,
num_images: int = 1,
num_inference_steps: int = 30,
guidance_scale: float = 7.5,
width: int = 512,
height: int = 512,
seed: Optional[int] = None
) -> List[Image.Image]:
"""
이미지 생성
Args:
prompt: 텍스트 프롬프트
negative_prompt: 부정 프롬프트 (원하지 않는 요소)
num_images: 생성할 이미지 개수
num_inference_steps: 디노이징 스텝 수 (높을수록 품질↑, 시간↑)
guidance_scale: 프롬프트 가이던스 (높을수록 프롬프트에 충실)
width, height: 이미지 크기 (8의 배수여야 함)
seed: 랜덤 시드 (재현성)
Returns:
생성된 이미지 리스트
"""
print(f"🎨 이미지 생성 시작")
print(f" 프롬프트: {prompt}")
print(f" 크기: {width}x{height}")
print(f" 스텝: {num_inference_steps}")
# 시드 설정
generator = None
if seed is not None:
generator = torch.Generator(device=self.device).manual_seed(seed)
print(f" 시드: {seed}")
# 기본 negative prompt
if negative_prompt is None:
negative_prompt = "blurry, bad quality, watermark, text, ugly, distorted"
# 생성
with torch.autocast(self.device):
result = self.pipe(
prompt=prompt,
negative_prompt=negative_prompt,
num_images_per_prompt=num_images,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
width=width,
height=height,
generator=generator
)
images = result.images
print(f"✅ {len(images)}개 이미지 생성 완료")
return images
def save_images(self, images: List[Image.Image], output_dir: str = "outputs"):
"""이미지 저장"""
os.makedirs(output_dir, exist_ok=True)
saved_paths = []
for i, img in enumerate(images):
filename = f"generated_{i+1}_{os.urandom(4).hex()}.png"
path = os.path.join(output_dir, filename)
img.save(path)
saved_paths.append(path)
print(f"💾 저장: {path}")
return saved_paths
# 테스트
if __name__ == "__main__":
# 생성기 초기화
generator = TextToImageGenerator(model_type="sd15", use_gpu=True)
# 프롬프트
prompts = [
"a beautiful sunset over mountains, landscape photography, highly detailed, 4k",
"a cute cat wearing a wizard hat, digital art, trending on artstation",
"cyberpunk city at night, neon lights, rain, cinematic lighting"
]
for prompt in prompts:
print(f"\n{'='*70}")
# 이미지 생성
images = generator.generate(
prompt=prompt,
num_images=2,
num_inference_steps=30,
guidance_scale=7.5,
width=512,
height=512,
seed=42 # 재현성을 위한 고정 시드
)
# 저장
paths = generator.save_images(images)
# 첫 번째 이미지 표시 (Jupyter/Colab에서)
# display(images[0])
🎨 Step 2: 프롬프트 엔지니어링 & 개선 (20분)
src/prompt_enhancer.py
import os
from openai import OpenAI
from dotenv import load_dotenv
from typing import Dict
load_dotenv()
class PromptEnhancer:
"""프롬프트 개선 및 최적화"""
def __init__(self):
self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
def enhance_prompt(self, simple_prompt: str) -> Dict[str, str]:
"""간단한 프롬프트를 상세하게 개선"""
system_prompt = """당신은 Stable Diffusion 프롬프트 엔지니어입니다.
사용자의 간단한 설명을 고품질 이미지 생성을 위한 상세한 프롬프트로 변환하세요.
좋은 프롬프트 작성 규칙:
1. 주제를 명확히 서술
2. 스타일/아트 형식 명시 (digital art, oil painting, photograph 등)
3. 분위기/조명 설명 (dramatic lighting, soft light, golden hour 등)
4. 품질 키워드 추가 (highly detailed, 4k, masterpiece, trending on artstation)
5. 구체적인 시각적 요소 포함 (colors, composition, perspective)
부정 프롬프트도 생성하세요 (원하지 않는 요소)."""
user_prompt = f"""간단한 설명: {simple_prompt}
다음 형식으로 답변하세요:
POSITIVE_PROMPT: [상세한 긍정 프롬프트]
NEGATIVE_PROMPT: [부정 프롬프트]"""
response = self.client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt}
],
temperature=0.7,
max_tokens=500
)
text = response.choices[0].message.content
# 파싱
positive = ""
negative = ""
if "POSITIVE_PROMPT:" in text:
positive = text.split("POSITIVE_PROMPT:")[1].split("NEGATIVE_PROMPT:")[0].strip()
if "NEGATIVE_PROMPT:" in text:
negative = text.split("NEGATIVE_PROMPT:")[1].strip()
return {
'original': simple_prompt,
'enhanced': positive,
'negative': negative
}
def generate_variations(self, base_prompt: str, num_variations: int = 3) -> list:
"""프롬프트 변형 생성"""
prompt = f"""다음 프롬프트의 {num_variations}가지 창의적 변형을 생성하세요.
각 변형은 다른 스타일이나 분위기를 가져야 합니다.
원본: {base_prompt}
변형 1:
변형 2:
변형 3:"""
response = self.client.chat.completions.create(
model="gpt-4o-mini",
messages=[{"role": "user", "content": prompt}],
temperature=0.8,
max_tokens=500
)
text = response.choices[0].message.content
variations = []
for line in text.split('\n'):
if line.strip().startswith('변형'):
var = line.split(':', 1)[1].strip() if ':' in line else line
variations.append(var)
return variations
# 테스트
if __name__ == "__main__":
enhancer = PromptEnhancer()
# 간단한 프롬프트들
simple_prompts = [
"고양이",
"산 풍경",
"미래 도시"
]
for simple in simple_prompts:
print(f"\n{'='*70}")
print(f"원본: {simple}")
print(f"{'='*70}")
result = enhancer.enhance_prompt(simple)
print(f"\n개선된 프롬프트:")
print(result['enhanced'])
print(f"\n부정 프롬프트:")
print(result['negative'])
# 변형 생성
print(f"\n프롬프트 변형:")
variations = enhancer.generate_variations(result['enhanced'], 3)
for i, var in enumerate(variations, 1):
print(f"{i}. {var}")
🖼️ Step 3: Image-to-Image 변환 (20분)
src/image_to_image.py
import torch
from PIL import Image
from diffusers import StableDiffusionImg2ImgPipeline
from typing import Optional
import sys
sys.path.append('..')
from utils.model_loader import ModelLoader
class ImageToImageGenerator:
"""Image-to-Image 변환"""
def __init__(self, use_gpu: bool = True):
loader = ModelLoader(use_gpu=use_gpu)
self.device = loader.device
print("📥 Image-to-Image 모델 로딩...")
self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
safety_checker=None
)
self.pipe = self.pipe.to(self.device)
if self.device == "cuda":
self.pipe.enable_attention_slicing()
print("✅ 모델 로드 완료")
def transform(
self,
init_image: Image.Image,
prompt: str,
strength: float = 0.75,
guidance_scale: float = 7.5,
num_inference_steps: int = 50,
seed: Optional[int] = None
) -> Image.Image:
"""
이미지 변환
Args:
init_image: 입력 이미지
prompt: 변환 프롬프트
strength: 변환 강도 (0.0=원본 유지, 1.0=완전 변환)
guidance_scale: 프롬프트 가이던스
num_inference_steps: 스텝 수
seed: 랜덤 시드
"""
print(f"🔄 이미지 변환 시작")
print(f" 프롬프트: {prompt}")
print(f" 강도: {strength}")
# 이미지 리사이즈 (512x512 권장)
init_image = init_image.resize((512, 512))
# 시드 설정
generator = None
if seed is not None:
generator = torch.Generator(device=self.device).manual_seed(seed)
# 변환
with torch.autocast(self.device):
result = self.pipe(
prompt=prompt,
image=init_image,
strength=strength,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
generator=generator
)
output_image = result.images[0]
print("✅ 변환 완료")
return output_image
# 테스트
if __name__ == "__main__":
from PIL import Image, ImageDraw
import os
# 테스트용 간단한 이미지 생성
def create_test_image():
img = Image.new('RGB', (512, 512), color='white')
draw = ImageDraw.Draw(img)
# 간단한 원 그리기
draw.ellipse([100, 100, 400, 400], fill='blue', outline='black', width=5)
os.makedirs("outputs", exist_ok=True)
img.save("outputs/test_input.png")
return img
# 생성기 초기화
generator = ImageToImageGenerator(use_gpu=True)
# 테스트 이미지
init_image = create_test_image()
# 변환 예시
transformations = [
("a blue planet in space, digital art", 0.5),
("a crystal ball, glass material, reflections", 0.7),
("abstract art, colorful, geometric patterns", 0.8)
]
for prompt, strength in transformations:
print(f"\n{'='*70}")
output = generator.transform(
init_image=init_image,
prompt=prompt,
strength=strength,
num_inference_steps=50,
seed=42
)
# 저장
filename = f"outputs/i2i_{strength}_{os.urandom(4).hex()}.png"
output.save(filename)
print(f"💾 저장: {filename}")
🎭 Step 4: Inpainting (이미지 편집) (25분)
src/inpainting.py
import torch
from PIL import Image, ImageDraw
from diffusers import StableDiffusionInpaintPipeline
from typing import Optional
import sys
sys.path.append('..')
from utils.model_loader import ModelLoader
class InpaintingGenerator:
"""Inpainting (이미지 부분 편집)"""
def __init__(self, use_gpu: bool = True):
loader = ModelLoader(use_gpu=use_gpu)
self.device = loader.device
print("📥 Inpainting 모델 로딩...")
self.pipe = StableDiffusionInpaintPipeline.from_pretrained(
"runwayml/stable-diffusion-inpainting",
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
safety_checker=None
)
self.pipe = self.pipe.to(self.device)
if self.device == "cuda":
self.pipe.enable_attention_slicing()
print("✅ 모델 로드 완료")
def inpaint(
self,
image: Image.Image,
mask: Image.Image,
prompt: str,
negative_prompt: Optional[str] = None,
num_inference_steps: int = 50,
guidance_scale: float = 7.5,
seed: Optional[int] = None
) -> Image.Image:
"""
마스크 영역 인페인팅
Args:
image: 원본 이미지
mask: 마스크 이미지 (흰색=편집 영역, 검은색=유지 영역)
prompt: 생성할 내용
negative_prompt: 부정 프롬프트
num_inference_steps: 스텝 수
guidance_scale: 가이던스
seed: 시드
"""
print(f"✂️ Inpainting 시작")
print(f" 프롬프트: {prompt}")
# 이미지 리사이즈
image = image.resize((512, 512))
mask = mask.resize((512, 512))
# 시드 설정
generator = None
if seed is not None:
generator = torch.Generator(device=self.device).manual_seed(seed)
# 기본 negative prompt
if negative_prompt is None:
negative_prompt = "blurry, bad quality, distorted"
# 인페인팅
with torch.autocast(self.device):
result = self.pipe(
prompt=prompt,
negative_prompt=negative_prompt,
image=image,
mask_image=mask,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
generator=generator
)
output_image = result.images[0]
print("✅ Inpainting 완료")
return output_image
def create_circular_mask(self, image: Image.Image, center_x: int, center_y: int, radius: int) -> Image.Image:
"""원형 마스크 생성"""
mask = Image.new('RGB', image.size, color='black')
draw = ImageDraw.Draw(mask)
# 흰색 원 그리기 (편집할 영역)
draw.ellipse(
[center_x - radius, center_y - radius, center_x + radius, center_y + radius],
fill='white'
)
return mask
def create_rectangle_mask(self, image: Image.Image, x1: int, y1: int, x2: int, y2: int) -> Image.Image:
"""사각형 마스크 생성"""
mask = Image.new('RGB', image.size, color='black')
draw = ImageDraw.Draw(mask)
# 흰색 사각형 그리기
draw.rectangle([x1, y1, x2, y2], fill='white')
return mask
# 테스트
if __name__ == "__main__":
from PIL import Image, ImageDraw
import os
# 테스트용 이미지 생성
def create_test_scene():
"""테스트 장면 생성"""
img = Image.new('RGB', (512, 512), color='skyblue')
draw = ImageDraw.Draw(img)
# 땅
draw.rectangle([0, 350, 512, 512], fill='green')
# 집
draw.rectangle([150, 250, 350, 350], fill='brown')
# 지붕
draw.polygon([150, 250, 250, 150, 350, 250], fill='red')
# 문
draw.rectangle([220, 280, 280, 350], fill='yellow')
os.makedirs("outputs", exist_ok=True)
img.save("outputs/test_scene.png")
return img
# 생성기 초기화
generator = InpaintingGenerator(use_gpu=True)
# 테스트 이미지
image = create_test_scene()
# 예시 1: 문을 나무 문으로 교체
print("\n" + "="*70)
print("예시 1: 문 교체")
mask1 = generator.create_rectangle_mask(image, 220, 280, 280, 350)
mask1.save("outputs/mask_door.png")
result1 = generator.inpaint(
image=image,
mask=mask1,
prompt="wooden door with metal handle, realistic, detailed",
num_inference_steps=50,
seed=42
)
result1.save("outputs/inpaint_door.png")
print("💾 저장: outputs/inpaint_door.png")
# 예시 2: 하늘에 해 추가
print("\n" + "="*70)
print("예시 2: 해 추가")
mask2 = generator.create_circular_mask(image, 400, 100, 60)
mask2.save("outputs/mask_sun.png")
result2 = generator.inpaint(
image=image,
mask=mask2,
prompt="bright sun, golden hour, warm lighting",
num_inference_steps=50,
seed=42
)
result2.save("outputs/inpaint_sun.png")
print("💾 저장: outputs/inpaint_sun.png")
🎮 Step 5: ControlNet (이미지 제어) (30분)
src/controlnet_generator.py
import torch
from PIL import Image
import numpy as np
import cv2
from diffusers import (
StableDiffusionControlNetPipeline,
ControlNetModel,
UniPCMultistepScheduler
)
from typing import Optional
import sys
sys.path.append('..')
from utils.model_loader import ModelLoader
class ControlNetGenerator:
"""ControlNet을 사용한 조건부 이미지 생성"""
def __init__(self, controlnet_type: str = "canny", use_gpu: bool = True):
"""
Args:
controlnet_type: 'canny', 'pose', 'depth', 'scribble' 등
"""
loader = ModelLoader(use_gpu=use_gpu)
self.device = loader.device
self.controlnet_type = controlnet_type
print(f"📥 ControlNet ({controlnet_type}) 로딩...")
# ControlNet 모델 선택
controlnet_models = {
'canny': 'lllyasviel/sd-controlnet-canny',
'pose': 'lllyasviel/sd-controlnet-openpose',
'depth': 'lllyasviel/sd-controlnet-depth',
'scribble': 'lllyasviel/sd-controlnet-scribble'
}
controlnet_model_id = controlnet_models.get(controlnet_type)
if not controlnet_model_id:
raise ValueError(f"지원하지 않는 ControlNet 타입: {controlnet_type}")
# ControlNet 로드
controlnet = ControlNetModel.from_pretrained(
controlnet_model_id,
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
)
# 파이프라인 생성
self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
controlnet=controlnet,
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
safety_checker=None
)
self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
self.pipe = self.pipe.to(self.device)
if self.device == "cuda":
self.pipe.enable_attention_slicing()
print("✅ ControlNet 로드 완료")
def preprocess_canny(self, image: Image.Image, low_threshold: int = 100, high_threshold: int = 200) -> Image.Image:
"""Canny 엣지 검출"""
image = np.array(image)
# BGR로 변환 (OpenCV 형식)
if len(image.shape) == 3 and image.shape[2] == 3:
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
# Canny 엣지 검출
edges = cv2.Canny(image, low_threshold, high_threshold)
# RGB로 변환
edges = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
return Image.fromarray(edges)
def generate_with_control(
self,
control_image: Image.Image,
prompt: str,
negative_prompt: Optional[str] = None,
num_inference_steps: int = 20,
guidance_scale: float = 9.0,
controlnet_conditioning_scale: float = 1.0,
seed: Optional[int] = None
) -> Image.Image:
"""
ControlNet을 사용한 이미지 생성
Args:
control_image: 제어 이미지 (Canny edge, pose 등)
prompt: 생성 프롬프트
negative_prompt: 부정 프롬프트
num_inference_steps: 스텝 수
guidance_scale: 가이던스
controlnet_conditioning_scale: ControlNet 강도 (0.0~2.0)
seed: 시드
"""
print(f"🎮 ControlNet 생성 시작")
print(f" 타입: {self.controlnet_type}")
print(f" 프롬프트: {prompt}")
# 이미지 리사이즈
control_image = control_image.resize((512, 512))
# 시드 설정
generator = None
if seed is not None:
generator = torch.Generator(device=self.device).manual_seed(seed)
# 기본 negative prompt
if negative_prompt is None:
negative_prompt = "blurry, bad quality, distorted"
# 생성
with torch.autocast(self.device):
result = self.pipe(
prompt=prompt,
negative_prompt=negative_prompt,
image=control_image,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
controlnet_conditioning_scale=controlnet_conditioning_scale,
generator=generator
)
output_image = result.images[0]
print("✅ 생성 완료")
return output_image
# 테스트
if __name__ == "__main__":
from PIL import Image
import os
# Canny ControlNet 테스트
generator = ControlNetGenerator(controlnet_type="canny", use_gpu=True)
# 테스트 이미지 생성 (간단한 도형)
def create_test_image():
img = Image.new('RGB', (512, 512), color='white')
from PIL import ImageDraw
draw = ImageDraw.Draw(img)
# 집 모양
draw.rectangle([100, 250, 400, 450], outline='black', width=5)
draw.polygon([100, 250, 250, 100, 400, 250], outline='black', width=5)
draw.rectangle([200, 350, 300, 450], outline='black', width=5)
os.makedirs("outputs", exist_ok=True)
img.save("outputs/test_drawing.png")
return img
# 테스트 이미지
test_image = create_test_image()
# Canny 엣지 추출
canny_image = generator.preprocess_canny(test_image)
canny_image.save("outputs/canny_edges.png")
print("💾 Canny 엣지 저장: outputs/canny_edges.png")
# 다양한 프롬프트로 생성
prompts = [
"a beautiful house in the countryside, oil painting",
"a modern house, architectural photography, high quality",
"a cottage in the forest, fantasy art, detailed"
]
for i, prompt in enumerate(prompts, 1):
print(f"\n{'='*70}")
result = generator.generate_with_control(
control_image=canny_image,
prompt=prompt,
num_inference_steps=20,
guidance_scale=9.0,
controlnet_conditioning_scale=1.0,
seed=42
)
filename = f"outputs/controlnet_{i}.png"
result.save(filename)
print(f"💾 저장: {filename}")
🎓 Step 6: LoRA 학습 (선택, 고급) (30분)
src/lora_trainer.py
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from diffusers.loaders import AttnProcsLayers
from diffusers.models.attention_processor import LoRAAttnProcessor
from PIL import Image
import os
from typing import List
class LoRATrainer:
"""LoRA (Low-Rank Adaptation) 학습"""
def __init__(self, base_model: str = "runwayml/stable-diffusion-v1-5", use_gpu: bool = True):
"""
Args:
base_model: 기본 모델
use_gpu: GPU 사용
"""
self.device = "cuda" if use_gpu and torch.cuda.is_available() else "cpu"
print(f"📥 기본 모델 로딩: {base_model}")
self.pipe = StableDiffusionPipeline.from_pretrained(
base_model,
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
safety_checker=None
)
self.pipe = self.pipe.to(self.device)
print("✅ 모델 로드 완료")
def prepare_lora(self, rank: int = 4):
"""LoRA 준비"""
print(f"🔧 LoRA 초기화 (rank={rank})")
# UNet에 LoRA 적용
unet = self.pipe.unet
lora_attn_procs = {}
for name in unet.attn_processors.keys():
cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
if name.startswith("mid_block"):
hidden_size = unet.config.block_out_channels[-1]
elif name.startswith("up_blocks"):
block_id = int(name[len("up_blocks.")])
hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
elif name.startswith("down_blocks"):
block_id = int(name[len("down_blocks.")])
hidden_size = unet.config.block_out_channels[block_id]
lora_attn_procs[name] = LoRAAttnProcessor(
hidden_size=hidden_size,
cross_attention_dim=cross_attention_dim,
rank=rank
)
unet.set_attn_processor(lora_attn_procs)
# LoRA 파라미터만 학습 가능하도록
for param in self.pipe.unet.parameters():
param.requires_grad = False
lora_layers = AttnProcsLayers(unet.attn_processors)
for param in lora_layers.parameters():
param.requires_grad = True
print("✅ LoRA 초기화 완료")
return lora_layers
def save_lora(self, lora_layers, save_path: str):
"""LoRA 가중치 저장"""
os.makedirs(os.path.dirname(save_path), exist_ok=True)
torch.save(lora_layers.state_dict(), save_path)
print(f"💾 LoRA 저장: {save_path}")
def load_lora(self, lora_path: str):
"""저장된 LoRA 로드"""
print(f"📥 LoRA 로딩: {lora_path}")
# LoRA 준비
lora_layers = self.prepare_lora()
# 가중치 로드
lora_layers.load_state_dict(torch.load(lora_path, map_location=self.device))
print("✅ LoRA 로드 완료")
# 간단한 사용 예시 (실제 학습은 복잡함)
if __name__ == "__main__":
trainer = LoRATrainer(use_gpu=True)
# LoRA 준비
lora_layers = trainer.prepare_lora(rank=4)
print("""
LoRA 학습을 위해서는 다음이 필요합니다:
1. 학습 이미지 데이터셋 (특정 스타일이나 대상)
2. 학습 루프 (optimizer, loss function)
3. 충분한 GPU 메모리 (8GB 이상)
4. 시간 (수백~수천 스텝)
자세한 학습 코드는 Hugging Face의 공식 예제를 참고하세요:
https://github.com/huggingface/diffusers/tree/main/examples/dreambooth
""")
# LoRA 저장 (학습 후)
# trainer.save_lora(lora_layers, "data/lora/my_style.pt")
🎨 Step 7: Streamlit UI - 통합 이미지 생성 서비스 (35분)
app.py
import streamlit as st
import sys
sys.path.append('src')
from text_to_image import TextToImageGenerator
from image_to_image import ImageToImageGenerator
from inpainting import InpaintingGenerator
from controlnet_generator import ControlNetGenerator
from prompt_enhancer import PromptEnhancer
from PIL import Image, ImageDraw
import io
import time
import os
# 페이지 설정
st.set_page_config(
page_title="AI Image Generation Service",
page_icon="🎨",
layout="wide"
)
# CSS
st.markdown("""
<style>
.main-title {
font-size: 3rem;
font-weight: bold;
text-align: center;
background: linear-gradient(120deg, #f093fb 0%, #f5576c 100%);
-webkit-background-clip: text;
-webkit-text-fill-color: transparent;
}
.stTabs [data-baseweb="tab-list"] {
gap: 10px;
}
.stTabs [data-baseweb="tab"] {
height: 50px;
background-color: #f0f2f6;
border-radius: 5px;
padding: 10px 20px;
}
.stTabs [aria-selected="true"] {
background-color: #ff4b4b;
color: white;
}
</style>
""", unsafe_allow_html=True)
# 세션 상태 초기화
if 't2i_generator' not in st.session_state:
st.session_state.t2i_generator = None
if 'i2i_generator' not in st.session_state:
st.session_state.i2i_generator = None
if 'inpaint_generator' not in st.session_state:
st.session_state.inpaint_generator = None
if 'controlnet_generator' not in st.session_state:
st.session_state.controlnet_generator = None
if 'prompt_enhancer' not in st.session_state:
st.session_state.prompt_enhancer = PromptEnhancer()
if 'generated_images' not in st.session_state:
st.session_state.generated_images = []
# 헤더
st.markdown('<h1 class="main-title">🎨 AI Image Generation Service</h1>', unsafe_allow_html=True)
st.markdown('<p style="text-align: center; font-size: 1.2rem; color: #666;">Stable Diffusion 기반 이미지 생성 플랫폼</p>', unsafe_allow_html=True)
# 사이드바
with st.sidebar:
st.header("⚙️ 시스템 설정")
use_gpu = st.checkbox(
"GPU 사용",
value=True,
help="GPU가 있으면 활성화하세요 (훨씬 빠름)"
)
st.markdown("---")
st.subheader("🖼️ 생성 기록")
if st.session_state.generated_images:
for i, img_info in enumerate(reversed(st.session_state.generated_images[-5:]), 1):
with st.expander(f"{i}. {img_info['type']}"):
st.image(img_info['image'], use_container_width=True)
st.caption(f"프롬프트: {img_info['prompt'][:50]}...")
st.caption(f"⏱️ {img_info['time']}")
else:
st.info("아직 생성한 이미지가 없습니다")
if st.button("🗑️ 기록 삭제"):
st.session_state.generated_images = []
st.rerun()
# 탭
tab1, tab2, tab3, tab4 = st.tabs([
"📝 Text-to-Image",
"🖼️ Image-to-Image",
"✂️ Inpainting",
"🎮 ControlNet"
])
# Tab 1: Text-to-Image
with tab1:
st.header("📝 Text-to-Image 생성")
col1, col2 = st.columns([2, 1])
with col1:
prompt = st.text_area(
"프롬프트",
height=100,
placeholder="예: a beautiful sunset over mountains, highly detailed, 4k"
)
# 프롬프트 개선
if st.button("✨ 프롬프트 개선"):
with st.spinner("프롬프트 개선 중..."):
enhanced = st.session_state.prompt_enhancer.enhance_prompt(prompt)
st.session_state.enhanced_prompt = enhanced['enhanced']
st.session_state.enhanced_negative = enhanced['negative']
st.success("프롬프트가 개선되었습니다!")
# 개선된 프롬프트 표시
if 'enhanced_prompt' in st.session_state:
with st.expander("개선된 프롬프트", expanded=True):
st.text_area("긍정 프롬프트", st.session_state.enhanced_prompt, height=100, key="show_enhanced")
st.text_area("부정 프롬프트", st.session_state.enhanced_negative, height=50, key="show_negative")
if st.button("개선된 프롬프트 사용"):
prompt = st.session_state.enhanced_prompt
negative_prompt = st.text_area(
"Negative Prompt",
value="blurry, bad quality, distorted",
height=50
)
with col2:
num_images = st.number_input("생성 개수", 1, 4, 1)
num_steps = st.slider("추론 스텝", 10, 50, 30)
guidance_scale = st.slider("Guidance Scale", 1.0, 20.0, 7.5)
col_w, col_h = st.columns(2)
width = col_w.selectbox("너비", [512, 768, 1024], index=0)
height = col_h.selectbox("높이", [512, 768, 1024], index=0)
seed = st.number_input("시드 (-1=랜덤)", -1, 999999, -1)
if seed == -1:
seed = None
if st.button("🎨 이미지 생성", type="primary", use_container_width=True):
if not prompt:
st.warning("프롬프트를 입력하세요!")
else:
# 생성기 초기화
if st.session_state.t2i_generator is None:
with st.spinner("모델 로딩 중... (최초 1회, 시간 소요)"):
st.session_state.t2i_generator = TextToImageGenerator(
model_type="sd15",
use_gpu=use_gpu
)
# 생성
with st.spinner(f"{num_images}개 이미지 생성 중..."):
start_time = time.time()
images = st.session_state.t2i_generator.generate(
prompt=prompt,
negative_prompt=negative_prompt,
num_images=num_images,
num_inference_steps=num_steps,
guidance_scale=guidance_scale,
width=width,
height=height,
seed=seed
)
elapsed = time.time() - start_time
# 결과 표시
st.success(f"✅ 생성 완료! ({elapsed:.1f}초)")
cols = st.columns(num_images)
for i, (col, img) in enumerate(zip(cols, images)):
with col:
st.image(img, use_container_width=True)
# 다운로드 버튼
buf = io.BytesIO()
img.save(buf, format='PNG')
st.download_button(
label="💾 다운로드",
data=buf.getvalue(),
file_name=f"generated_{i+1}.png",
mime="image/png",
key=f"download_t2i_{i}"
)
# 기록 저장
st.session_state.generated_images.append({
'type': 'Text-to-Image',
'image': img,
'prompt': prompt,
'time': time.strftime('%H:%M:%S')
})
# Tab 2: Image-to-Image
with tab2:
st.header("🖼️ Image-to-Image 변환")
col1, col2 = st.columns([2, 1])
with col1:
uploaded_file = st.file_uploader("이미지 업로드", type=['png', 'jpg', 'jpeg'])
if uploaded_file:
input_image = Image.open(uploaded_file)
st.image(input_image, caption="입력 이미지", use_container_width=True)
i2i_prompt = st.text_area(
"변환 프롬프트",
height=100,
placeholder="예: oil painting style, vibrant colors"
)
with col2:
strength = st.slider("변환 강도", 0.0, 1.0, 0.75, help="높을수록 원본에서 멀어짐")
i2i_steps = st.slider("스텝 수", 10, 100, 50)
i2i_guidance = st.slider("가이던스", 1.0, 20.0, 7.5)
i2i_seed = st.number_input("시드", -1, 999999, -1, key="i2i_seed")
if i2i_seed == -1:
i2i_seed = None
if st.button("🔄 변환하기", type="primary", key="i2i_button"):
if not uploaded_file:
st.warning("이미지를 업로드하세요!")
elif not i2i_prompt:
st.warning("프롬프트를 입력하세요!")
else:
# 생성기 초기화
if st.session_state.i2i_generator is None:
with st.spinner("모델 로딩 중..."):
st.session_state.i2i_generator = ImageToImageGenerator(use_gpu=use_gpu)
# 변환
with st.spinner("변환 중..."):
output = st.session_state.i2i_generator.transform(
init_image=input_image,
prompt=i2i_prompt,
strength=strength,
guidance_scale=i2i_guidance,
num_inference_steps=i2i_steps,
seed=i2i_seed
)
# 결과
st.success("✅ 변환 완료!")
col_before, col_after = st.columns(2)
with col_before:
st.image(input_image, caption="변환 전", use_container_width=True)
with col_after:
st.image(output, caption="변환 후", use_container_width=True)
# 다운로드
buf = io.BytesIO()
output.save(buf, format='PNG')
st.download_button(
"💾 다운로드",
buf.getvalue(),
"transformed.png",
"image/png"
)
# Tab 3: Inpainting
with tab3:
st.header("✂️ Inpainting (부분 편집)")
st.info("💡 팁: 마스크 도구를 사용하여 편집할 영역을 지정하세요")
col1, col2 = st.columns([2, 1])
with col1:
inpaint_file = st.file_uploader("편집할 이미지", type=['png', 'jpg', 'jpeg'], key="inpaint_upload")
if inpaint_file:
inpaint_image = Image.open(inpaint_file)
st.image(inpaint_image, caption="원본", use_container_width=True)
# 간단한 마스크 생성 옵션
mask_type = st.radio(
"마스크 타입",
["원형", "사각형", "이미지 업로드"]
)
if mask_type == "원형":
col_x, col_y, col_r = st.columns(3)
mask_x = col_x.number_input("중심 X", 0, 512, 256)
mask_y = col_y.number_input("중심 Y", 0, 512, 256)
mask_r = col_r.number_input("반지름", 10, 256, 100)
elif mask_type == "사각형":
col_x1, col_y1 = st.columns(2)
col_x2, col_y2 = st.columns(2)
mask_x1 = col_x1.number_input("좌측 상단 X", 0, 512, 100)
mask_y1 = col_y1.number_input("좌측 상단 Y", 0, 512, 100)
mask_x2 = col_x2.number_input("우측 하단 X", 0, 512, 400)
mask_y2 = col_y2.number_input("우측 하단 Y", 0, 512, 400)
else: # 이미지 업로드
mask_file = st.file_uploader("마스크 이미지 (흰색=편집)", type=['png'], key="mask_upload")
inpaint_prompt = st.text_area(
"생성할 내용",
placeholder="예: a red rose, detailed, realistic",
key="inpaint_prompt"
)
with col2:
inpaint_steps = st.slider("스텝", 20, 100, 50, key="inpaint_steps")
inpaint_guidance = st.slider("가이던스", 1.0, 20.0, 7.5, key="inpaint_guidance")
inpaint_seed = st.number_input("시드", -1, 999999, -1, key="inpaint_seed_input")
if inpaint_seed == -1:
inpaint_seed = None
if st.button("✂️ Inpaint", type="primary", key="inpaint_button"):
if not inpaint_file:
st.warning("이미지를 업로드하세요!")
elif not inpaint_prompt:
st.warning("프롬프트를 입력하세요!")
else:
# 생성기 초기화
if st.session_state.inpaint_generator is None:
with st.spinner("모델 로딩 중..."):
st.session_state.inpaint_generator = InpaintingGenerator(use_gpu=use_gpu)
# 마스크 생성
if mask_type == "원형":
mask = st.session_state.inpaint_generator.create_circular_mask(
inpaint_image, mask_x, mask_y, mask_r
)
elif mask_type == "사각형":
mask = st.session_state.inpaint_generator.create_rectangle_mask(
inpaint_image, mask_x1, mask_y1, mask_x2, mask_y2
)
else:
if mask_file:
mask = Image.open(mask_file)
else:
st.error("마스크 이미지를 업로드하세요!")
st.stop()
# Inpaint
with st.spinner("Inpainting 중..."):
result = st.session_state.inpaint_generator.inpaint(
image=inpaint_image,
mask=mask,
prompt=inpaint_prompt,
num_inference_steps=inpaint_steps,
guidance_scale=inpaint_guidance,
seed=inpaint_seed
)
st.success("✅ Inpainting 완료!")
col_orig, col_mask, col_result = st.columns(3)
with col_orig:
st.image(inpaint_image, caption="원본", use_container_width=True)
with col_mask:
st.image(mask, caption="마스크", use_container_width=True)
with col_result:
st.image(result, caption="결과", use_container_width=True)
# Tab 4: ControlNet
with tab4:
st.header("🎮 ControlNet (구조 제어)")
st.info("💡 이미지의 구조(윤곽선)를 유지하면서 스타일을 변경합니다")
col1, col2 = st.columns([2, 1])
with col1:
control_file = st.file_uploader(
"입력 이미지 (구조 추출용)",
type=['png', 'jpg', 'jpeg'],
key="control_upload"
)
if control_file:
control_image = Image.open(control_file)
st.image(control_image, caption="입력", use_container_width=True)
control_prompt = st.text_area(
"생성 프롬프트",
placeholder="예: oil painting of a house, detailed, artistic",
key="control_prompt"
)
with col2:
controlnet_type = st.selectbox(
"ControlNet 타입",
["canny"] # 다른 타입 추가 가능
)
if controlnet_type == "canny":
low_threshold = st.slider("Canny Low", 50, 150, 100)
high_threshold = st.slider("Canny High", 150, 250, 200)
control_steps = st.slider("스텝", 10, 50, 20, key="control_steps")
control_guidance = st.slider("가이던스", 1.0, 20.0, 9.0, key="control_guidance")
control_scale = st.slider("제어 강도", 0.0, 2.0, 1.0, key="control_scale")
control_seed = st.number_input("시드", -1, 999999, -1, key="control_seed_input")
if control_seed == -1:
control_seed = None
if st.button("🎮 생성", type="primary", key="control_button"):
if not control_file:
st.warning("이미지를 업로드하세요!")
elif not control_prompt:
st.warning("프롬프트를 입력하세요!")
else:
# 생성기 초기화
if st.session_state.controlnet_generator is None:
with st.spinner("ControlNet 로딩 중... (시간 소요)"):
st.session_state.controlnet_generator = ControlNetGenerator(
controlnet_type=controlnet_type,
use_gpu=use_gpu
)
# 전처리
with st.spinner("엣지 추출 중..."):
if controlnet_type == "canny":
control_processed = st.session_state.controlnet_generator.preprocess_canny(
control_image,
low_threshold,
high_threshold
)
# 생성
with st.spinner("ControlNet 생성 중..."):
result = st.session_state.controlnet_generator.generate_with_control(
control_image=control_processed,
prompt=control_prompt,
num_inference_steps=control_steps,
guidance_scale=control_guidance,
controlnet_conditioning_scale=control_scale,
seed=control_seed
)
st.success("✅ 생성 완료!")
col_orig, col_edge, col_result = st.columns(3)
with col_orig:
st.image(control_image, caption="원본", use_container_width=True)
with col_edge:
st.image(control_processed, caption="엣지", use_container_width=True)
with col_result:
st.image(result, caption="생성 결과", use_container_width=True)
# Footer
st.markdown("---")
st.markdown("""
<div style='text-align: center; color: #666;'>
<p><strong>Stable Diffusion Image Generation Service</strong></p>
<p>Text-to-Image • Image-to-Image • Inpainting • ControlNet</p>
</div>
""", unsafe_allow_html=True)
실행
streamlit run app.py
🎓 학습 과제 및 다음 단계
필수 과제 체크리스트
- [ ] Text-to-Image 생성 성공
- [ ] 프롬프트 개선 기능 테스트
- [ ] Image-to-Image 변환 확인
- [ ] Inpainting 동작 확인
- [ ] ControlNet (Canny) 테스트
심화 과제
- LoRA 학습: 특정 스타일이나 캐릭터 학습
- 배치 처리: 여러 이미지 한 번에 생성
- API 서버: FastAPI로 REST API 구축
- 이미지 품질 향상: Upscaling, Face restoration
- 비디오 생성: AnimateDiff로 동영상 생성
프로덕션 배포
- Docker 컨테이너화
- GPU 서버 최적화
- 큐 시스템 (Celery + Redis)
- 캐싱 전략
- 비용 관리
반응형