import os
import gc
import re
import gradio as gr
import numpy as np
import torch
import json
import spaces
import config
import utils
import logging
import time
from datetime import datetime
from typing import List, Dict, Tuple, Optional
from PIL import Image, PngImagePlugin
from diffusers.models import AutoencoderKL
from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
from transformers import pipeline as translation_pipeline
from config import (
MODEL,
MIN_IMAGE_SIZE,
MAX_IMAGE_SIZE,
USE_TORCH_COMPILE,
ENABLE_CPU_OFFLOAD,
OUTPUT_DIR,
DEFAULT_NEGATIVE_PROMPT,
DEFAULT_ASPECT_RATIO,
examples,
sampler_list,
aspect_ratios,
style_list,
)
# Enhanced logging configuration
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)
# Constants
IS_COLAB = utils.is_google_colab() or os.getenv("IS_COLAB") == "1"
HF_TOKEN = os.getenv("HF_TOKEN")
CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"
# PyTorch settings for better performance and determinism
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.backends.cuda.matmul.allow_tf32 = True
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
logger.info(f"Using device: {device}")
# 번역 파이프라인 초기화 (한글 → 영어)
translator = translation_pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
class GenerationError(Exception):
"""Custom exception for generation errors"""
pass
def translate_if_korean(prompt: str) -> str:
"""프롬프트에 한글이 포함되어 있으면 영어로 번역"""
if re.search(r'[ㄱ-ㅎㅏ-ㅣ가-힣]', prompt):
logger.info("Korean detected in prompt. Translating to English...")
try:
translation = translator(prompt)[0]['translation_text']
logger.info(f"Translation result: {translation}")
return translation
except Exception as e:
logger.error(f"Translation error: {e}")
# 번역 실패 시 원본 프롬프트 사용
return prompt
return prompt
def validate_prompt(prompt: str) -> str:
"""Validate and clean up the input prompt."""
if not isinstance(prompt, str):
raise GenerationError("Prompt must be a string")
try:
# Ensure proper UTF-8 encoding/decoding
prompt = prompt.encode('utf-8').decode('utf-8')
# Add space between ! and ,
prompt = prompt.replace("!,", "! ,")
except UnicodeError:
raise GenerationError("Invalid characters in prompt")
# Only check if the prompt is completely empty or only whitespace
if not prompt or prompt.isspace():
raise GenerationError("Prompt cannot be empty")
# 번역 적용: 한글이 감지되면 영어로 변환
prompt = translate_if_korean(prompt)
return prompt.strip()
def validate_dimensions(width: int, height: int) -> None:
"""Validate image dimensions."""
if not MIN_IMAGE_SIZE <= width <= MAX_IMAGE_SIZE:
raise GenerationError(f"Width must be between {MIN_IMAGE_SIZE} and {MAX_IMAGE_SIZE}")
if not MIN_IMAGE_SIZE <= height <= MAX_IMAGE_SIZE:
raise GenerationError(f"Height must be between {MIN_IMAGE_SIZE} and {MAX_IMAGE_SIZE}")
@spaces.GPU
def generate(
prompt: str,
negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
seed: int = 0,
custom_width: int = 1024, # 기본 크기를 1024로 설정
custom_height: int = 1024, # 기본 크기를 1024로 설정
guidance_scale: float = 6.0,
num_inference_steps: int = 25,
sampler: str = "Euler a",
aspect_ratio_selector: str = DEFAULT_ASPECT_RATIO,
style_selector: str = "(None)",
use_upscaler: bool = False,
upscaler_strength: float = 0.55,
upscale_by: float = 1.5,
add_quality_tags: bool = True,
progress: gr.Progress = gr.Progress(track_tqdm=True),
) -> Tuple[List[str], Dict]:
"""Generate images based on the given parameters."""
start_time = time.time()
upscaler_pipe = None
backup_scheduler = None
try:
# Memory management
torch.cuda.empty_cache()
gc.collect()
# Input validation
prompt = validate_prompt(prompt)
if negative_prompt:
negative_prompt = negative_prompt.encode('utf-8').decode('utf-8')
validate_dimensions(custom_width, custom_height)
# Set up generation
generator = utils.seed_everything(seed)
width, height = utils.aspect_ratio_handler(
aspect_ratio_selector,
custom_width,
custom_height,
)
# Process prompts
if add_quality_tags:
prompt = f"masterpiece, high score, great score, absurdres, {prompt}"
prompt, negative_prompt = utils.preprocess_prompt(
styles, style_selector, prompt, negative_prompt
)
width, height = utils.preprocess_image_dimensions(width, height)
# Set up pipeline
backup_scheduler = pipe.scheduler
pipe.scheduler = utils.get_scheduler(pipe.scheduler.config, sampler)
if use_upscaler:
upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
# Prepare metadata
metadata = {
"prompt": prompt,
"negative_prompt": negative_prompt,
"resolution": f"{width} x {height}",
"guidance_scale": guidance_scale,
"num_inference_steps": num_inference_steps,
"style_preset": style_selector,
"seed": seed,
"sampler": sampler,
"Model": "Animagine XL 4.0",
"Model hash": "e3c47aedb0",
}
if use_upscaler:
new_width = int(width * upscale_by)
new_height = int(height * upscale_by)
metadata["use_upscaler"] = {
"upscale_method": "nearest-exact",
"upscaler_strength": upscaler_strength,
"upscale_by": upscale_by,
"new_resolution": f"{new_width} x {new_height}",
}
else:
metadata["use_upscaler"] = None
logger.info(f"Starting generation with parameters: {json.dumps(metadata, indent=4)}")
# Generate images
if use_upscaler:
latents = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
width=width,
height=height,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
generator=generator,
output_type="latent",
).images
upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
images = upscaler_pipe(
prompt=prompt,
negative_prompt=negative_prompt,
image=upscaled_latents,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
strength=upscaler_strength,
generator=generator,
output_type="pil",
).images
else:
images = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
width=width,
height=height,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
generator=generator,
output_type="pil",
).images
# Save images
if images:
total = len(images)
image_paths = []
for idx, image in enumerate(images, 1):
progress(idx/total, desc="Saving images...")
path = utils.save_image(image, metadata, OUTPUT_DIR, IS_COLAB)
image_paths.append(path)
logger.info(f"Image {idx}/{total} saved as {path}")
generation_time = time.time() - start_time
logger.info(f"Generation completed successfully in {generation_time:.2f} seconds")
metadata["generation_time"] = f"{generation_time:.2f}s"
return image_paths, metadata
except GenerationError as e:
logger.warning(f"Generation validation error: {str(e)}")
raise gr.Error(str(e))
except Exception as e:
logger.exception("Unexpected error during generation")
raise gr.Error(f"Generation failed: {str(e)}")
finally:
# Cleanup
torch.cuda.empty_cache()
gc.collect()
if upscaler_pipe is not None:
del upscaler_pipe
if backup_scheduler is not None and pipe is not None:
pipe.scheduler = backup_scheduler
utils.free_memory()
# Model initialization
if torch.cuda.is_available():
try:
logger.info("Loading VAE and pipeline...")
vae = AutoencoderKL.from_pretrained(
"madebyollin/sdxl-vae-fp16-fix",
torch_dtype=torch.float16,
)
pipe = utils.load_pipeline(MODEL, device, vae=vae)
logger.info("Pipeline loaded successfully on GPU!")
except Exception as e:
logger.error(f"Error loading VAE, falling back to default: {e}")
pipe = utils.load_pipeline(MODEL, device)
else:
logger.warning("CUDA not available, running on CPU")
pipe = None
# Process styles
styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
# 사용자 인터페이스 (UI) 개선: CSS 스타일 및 레이아웃 수정
custom_css = """
/* 배경 및 글자 색상 변경 */
body {
background-color: #f7f9fc;
color: #333;
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
}
/* 헤더 스타일 */
.header {
text-align: center;
padding: 20px;
}
.header .title {
font-size: 3em;
font-weight: bold;
color: #2c3e50;
}
.header .subtitle {
font-size: 1.2em;
color: #7f8c8d;
}
a {
text-decoration: none;
color: #3498db;
}
/* Discord 버튼 스타일 */
.discord-btn {
display: flex;
align-items: center;
justify-content: center;
padding: 10px 20px;
background: #7289da;
color: white;
border-radius: 8px;
font-weight: bold;
margin-top: 20px;
}
.discord-btn:hover {
background: #5b6eae;
}
.discord-icon {
width: 24px;
height: 24px;
margin-right: 8px;
}
/* Gradio 갤러리 스타일 개선 */
.gradio-gallery {
border: none;
box-shadow: none;
}
"""
with gr.Blocks(css=custom_css, theme="default") as demo:
# 상단 헤더
gr.HTML(
"""
"""
)
# 좌측에 입력 및 설정, 우측에 결과 출력 배치
with gr.Row():
with gr.Column(scale=3):
with gr.Tabs():
with gr.TabItem("Generate"):
with gr.Group():
prompt = gr.Textbox(
label="Prompt",
lines=4,
placeholder="Describe what you want to generate...",
info="Enter your image generation prompt here. 한글 입력 시 자동으로 영어로 번역됩니다.",
)
negative_prompt = gr.Textbox(
label="Negative Prompt",
lines=4,
placeholder="Describe what you want to avoid",
value=DEFAULT_NEGATIVE_PROMPT,
info="Specify elements you don't want in the image.",
)
add_quality_tags = gr.Checkbox(
label="Quality Tags",
value=True,
info="Automatically add quality-enhancing tags to your prompt.",
)
with gr.Accordion(label="More Settings", open=False):
with gr.Column():
aspect_ratio_selector = gr.Radio(
label="Aspect Ratio",
choices=aspect_ratios,
value=DEFAULT_ASPECT_RATIO,
container=True,
info="Choose the dimensions of your image.",
)
with gr.Row(visible=False) as custom_resolution:
custom_width = gr.Slider(
label="Width",
minimum=MIN_IMAGE_SIZE,
maximum=MAX_IMAGE_SIZE,
step=8,
value=1024, # 기본값 1024
info=f"Image width (between {MIN_IMAGE_SIZE} and {MAX_IMAGE_SIZE})",
)
custom_height = gr.Slider(
label="Height",
minimum=MIN_IMAGE_SIZE,
maximum=MAX_IMAGE_SIZE,
step=8,
value=1024, # 기본값 1024
info=f"Image height (between {MIN_IMAGE_SIZE} and {MAX_IMAGE_SIZE})",
)
with gr.Accordion(label="Advanced Parameters", open=False):
with gr.Row():
style_selector = gr.Dropdown(
label="Style Preset",
choices=list(styles.keys()),
value="(None)",
info="Apply a predefined style to your generation.",
)
sampler = gr.Dropdown(
label="Sampler",
choices=sampler_list,
value="Euler a",
info="Different samplers can produce varying results.",
)
with gr.Row():
seed = gr.Slider(
label="Seed",
minimum=0,
maximum=utils.MAX_SEED,
step=1,
value=0,
info="Set a specific seed for reproducible results.",
)
randomize_seed = gr.Checkbox(
label="Randomize seed",
value=True,
info="Generate a new random seed for each image.",
)
with gr.Row():
guidance_scale = gr.Slider(
label="Guidance scale",
minimum=1,
maximum=12,
step=0.1,
value=6.0,
info="Higher values make the image more closely match your prompt.",
)
num_inference_steps = gr.Slider(
label="Inference steps",
minimum=1,
maximum=50,
step=1,
value=25,
info="More steps generally yield higher quality images.",
)
with gr.Row():
use_upscaler = gr.Checkbox(
label="Use Upscaler",
value=False,
info="Enable high-resolution upscaling.",
)
upscaler_strength = gr.Slider(
label="Upscaler Strength",
minimum=0,
maximum=1,
step=0.05,
value=0.55,
visible=False,
info="Control how strongly the upscaler affects the image.",
)
upscale_by = gr.Slider(
label="Upscale By",
minimum=1,
maximum=1.5,
step=0.1,
value=1.5,
visible=False,
info="Multiplier for the final image resolution.",
)
with gr.TabItem("Examples"):
gr.Markdown(
"""
### Example Prompts
- **Scenic Landscape:** A breathtaking view of a mountain landscape during sunrise.
- **Cyberpunk City:** A futuristic cyberpunk city with neon lights and towering skyscrapers.
- **Fantasy Character:** A majestic wizard with a long beard and glowing magical staff.
"""
)
gr.Examples(
examples=examples,
inputs=prompt,
outputs=[],
cache_examples=CACHE_EXAMPLES,
)
run_button = gr.Button("Generate", variant="primary", elem_id="generate-button")
with gr.Column(scale=4):
result = gr.Gallery(
label="Generated Images",
columns=2,
height="600px",
show_label=True,
elem_classes="gradio-gallery",
)
with gr.Accordion(label="Generation Parameters", open=False):
gr_metadata = gr.JSON(
label="Image Metadata",
show_label=True,
)
# Discord 버튼을 하단 중앙에 배치
with gr.Row():
gr.HTML(
"""
"""
)
# UI 동적 업데이트: 업스케일러 옵션 및 커스텀 해상도 슬라이더 표시 제어
use_upscaler.change(
fn=lambda x: [gr.update(visible=x), gr.update(visible=x)],
inputs=use_upscaler,
outputs=[upscaler_strength, upscale_by],
queue=False,
api_name=False,
)
aspect_ratio_selector.change(
fn=lambda x: gr.update(visible=x == "Custom"),
inputs=aspect_ratio_selector,
outputs=custom_resolution,
queue=False,
api_name=False,
)
# 트리거 설정: 프롬프트 입력, 버튼 클릭 시 시드 랜덤화 후 이미지 생성 실행
gr.on(
triggers=[
prompt.submit,
negative_prompt.submit,
run_button.click,
],
fn=utils.randomize_seed_fn,
inputs=[seed, randomize_seed],
outputs=seed,
queue=False,
api_name=False,
).then(
fn=lambda: gr.update(interactive=False, value="Generating..."),
outputs=run_button,
).then(
fn=generate,
inputs=[
prompt,
negative_prompt,
seed,
custom_width,
custom_height,
guidance_scale,
num_inference_steps,
sampler,
aspect_ratio_selector,
style_selector,
use_upscaler,
upscaler_strength,
upscale_by,
add_quality_tags,
],
outputs=[result, gr_metadata],
).then(
fn=lambda: gr.update(interactive=True, value="Generate"),
outputs=run_button,
)
if __name__ == "__main__":
demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)