Add initial project structure for AI Comic Generation

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.ttf filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+__pycache__
+weights/
+*.zip
+.env

app.py

@@ -0,0 +1,203 @@
+import os
+import base64
+import logging
+
+import gradio as gr
+from dotenv import load_dotenv
+
+from src.config import settings
+from src.llm import LLMClient
+from src.stable_diffusion import DiffusionClient
+from src.story_generator import StoryGenerator
+from src.image_prompt_generator import ImagePromptGenerator
+from src.image_generator import ImageGenerator
+
+load_dotenv()
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# ── Lazy-loaded singleton pipeline components ────────────────────────────────
+llm_client: LLMClient | None = None
+story_gen: StoryGenerator | None = None
+prompt_gen: ImagePromptGenerator | None = None
+img_gen: ImageGenerator | None = None
+
+
+def _init_pipeline():
+    """Instantiate and load all models once on first call."""
+    global llm_client, story_gen, prompt_gen, img_gen
+    if llm_client is not None:
+        return
+
+    logger.info("Initializing comic generation pipeline...")
+    llm_client = LLMClient(
+        hf_token=settings.HF_TOKEN,
+        cache_dir=settings.CACHE_DIR,
+    )
+    llm_client.load_model()
+
+    diff_client = DiffusionClient(
+        hf_token=settings.HF_TOKEN,
+        cache_dir=settings.CACHE_DIR,
+    )
+    diff_client.load_model()
+
+    story_gen = StoryGenerator(llm_client)
+    prompt_gen = ImagePromptGenerator(llm_client)
+    img_gen = ImageGenerator(diff_client)
+    logger.info("Pipeline ready.")
+
+
+# ── Core inference function ──────────────────────────────────────────────────
+def inference(user_prompt: str):
+    """Run the full comic generation pipeline and yield progress updates.
+
+    Returns:
+        Tuple of (gallery_images, story_text, status_message).
+    """
+    if not user_prompt or not user_prompt.strip():
+        return [], "Vui lòng nhập ý tưởng truyện.", "⚠️ Chưa nhập nội dung."
+
+    _init_pipeline()
+
+    yield [], "", "⏳ Đang tạo câu chuyện..."
+
+    raw_story = story_gen.gen_story_structured(user_prompt)
+    paragraphs = prompt_gen.parse_paragraphs(raw_story)
+
+    if not paragraphs:
+        yield [], raw_story, "⚠️ Không tách được đoạn văn. Xem raw output bên dưới."
+        return
+
+    max_scenes = min(len(paragraphs), settings.MAX_SCENES)
+    paragraphs = paragraphs[:max_scenes]
+
+    story_display = f"**Raw LLM output:**\n```\n{raw_story}\n```\n\n"
+    story_display += "**Parsed paragraphs:**\n"
+    for i, p in enumerate(paragraphs, 1):
+        story_display += f"{i}. {p}\n"
+
+    yield [], story_display, f"⏳ Đang tạo hình ảnh cho {max_scenes} cảnh..."
+
+    images = []
+    for idx, paragraph in enumerate(paragraphs):
+        visual_prompt = prompt_gen.create_visual_prompt(
+            paragraph, raw_story, settings.COMIC_STYLE
+        )
+        story_display += f"\n**Prompt {idx + 1}:** {visual_prompt}\n"
+        yield images, story_display, f"🎨 Đang vẽ cảnh {idx + 1}/{max_scenes}..."
+
+        panel = img_gen.generate_image(
+            prompt=visual_prompt,
+            paragraph=paragraph,
+        )
+        if panel is not None:
+            images.append(panel)
+            yield images, story_display, f"✅ Hoàn thành cảnh {idx + 1}/{max_scenes}"
+
+    yield images, story_display, f"🎉 Hoàn tất! Đã tạo {len(images)} panel truyện tranh."
+
+
+# ── UI helpers ───────────────────────────────────────────────────────────────
+def _image_to_base64(image_path: str) -> str:
+    with open(image_path, "rb") as f:
+        return base64.b64encode(f.read()).decode("utf-8")
+
+
+def _create_header():
+    with gr.Row():
+        with gr.Column(scale=1):
+            logo_b64 = _image_to_base64("static/aivn_logo.png")
+            gr.HTML(
+                f'<img src="data:image/png;base64,{logo_b64}" '
+                'alt="Logo" style="height:120px;width:auto;margin-right:20px;margin-bottom:20px;">'
+            )
+        with gr.Column(scale=4):
+            gr.Markdown(
+                """
+                <div style="display:flex;justify-content:space-between;align-items:center;padding:0 15px;">
+                    <div>
+                        <h1 style="margin-bottom:0;">🎨 AI Comic Generation Demo</h1>
+                        <p style="margin-top:0.5em;color:#666;">Tạo truyện tranh từ ý tưởng bằng AI (LLM + SDXL)</p>
+                    </div>
+                    <div style="text-align:right;border-left:2px solid #ddd;padding-left:20px;">
+                        <h3 style="margin:0;color:#2c3e50;">🚀 AIO2025 Project</h3>
+                        <p style="margin:0;color:#7f8c8d;">Comic Generation with LLM & Stable Diffusion XL</p>
+                    </div>
+                </div>
+                """
+            )
+
+
+def _create_footer():
+    gr.HTML(
+        """
+        <style>
+            .sticky-footer {
+                position:fixed;bottom:0;left:0;width:100%;
+                background:white;padding:10px;
+                box-shadow:0 -2px 10px rgba(0,0,0,0.1);z-index:1000;
+            }
+            .content-wrap { padding-bottom:60px; }
+        </style>
+        <div class="sticky-footer">
+            <div style="text-align:center;font-size:14px;">
+                Created by <a href="https://vlai.work" target="_blank"
+                style="color:#007BFF;text-decoration:none;">VLAI</a> • AI VIETNAM
+            </div>
+        </div>
+        """
+    )
+
+
+# ── App layout ───────────────────────────────────────────────────────────────
+custom_css = """
+.gradio-container { min-height:100vh; }
+.content-wrap { padding-bottom:60px; }
+.full-width-btn {
+    width:100% !important; height:50px !important;
+    font-size:18px !important; margin-top:20px !important;
+    background:linear-gradient(45deg,#FF6B6B,#4ECDC4) !important;
+    color:white !important; border:none !important;
+}
+.full-width-btn:hover {
+    background:linear-gradient(45deg,#FF5252,#3CB4AC) !important;
+}
+"""
+
+with gr.Blocks(css=custom_css) as demo:
+    _create_header()
+
+    with gr.Column(variant="panel"):
+        prompt_input = gr.Textbox(
+            label="💡 Tóm tắt truyện (Tiếng Việt)",
+            placeholder="Ví dụ: Một con thỏ lười biếng học được bài học về sự chăm chỉ từ một con rùa...",
+            lines=3,
+        )
+        generate_btn = gr.Button(
+            "Tạo Truyện Tranh 🎨",
+            elem_classes="full-width-btn",
+        )
+        status_box = gr.Textbox(label="Trạng thái", interactive=False)
+
+    with gr.Row(equal_height=True):
+        gallery = gr.Gallery(
+            label="🖼️ Truyện tranh",
+            columns=2,
+            object_fit="contain",
+            height="auto",
+        )
+
+    with gr.Accordion("📝 Chi tiết câu chuyện & prompts", open=False):
+        story_output = gr.Markdown()
+
+    generate_btn.click(
+        fn=inference,
+        inputs=prompt_input,
+        outputs=[gallery, story_output, status_box],
+    )
+
+    _create_footer()
+
+if __name__ == "__main__":
+    demo.launch(allowed_paths=["static/aivn_logo.png"])

requirements.txt

@@ -0,0 +1,11 @@
+torch==2.6.0
+diffusers==0.32.2
+transformers==4.50.3
+accelerate==1.6.0
+gradio==6.9.0
+numpy==1.26.4
+pillow
+pydantic-settings
+python-dotenv
+sentencepiece
+bitsandbytes>=0.46.1

src/config.py

@@ -0,0 +1,19 @@
+import os
+from dotenv import load_dotenv
+from pydantic_settings import BaseSettings
+
+load_dotenv()
+
+
+class Settings(BaseSettings):
+    HF_TOKEN: str = os.getenv("HF_TOKEN", "")
+    CACHE_DIR: str = os.getenv("CACHE_DIR", "./cache")
+    MAX_SCENES: int = 4
+    COMIC_STYLE: str = (
+        "whimsical watercolor children's book illustration, "
+        "hand-drawn, soft pencil outlines, pastel colors, "
+        "cozy atmosphere, high quality, magical lighting"
+    )
+
+
+settings = Settings()

src/image_generator.py

@@ -0,0 +1,65 @@
+import os
+import textwrap
+import logging
+from PIL import Image, ImageDraw, ImageFont, ImageOps
+from src.stable_diffusion import DiffusionClient
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+FONT_PATH = os.path.join(os.path.dirname(__file__), "..", "static", "Roboto-Regular.ttf")
+
+
+class ImageGenerator:
+    def __init__(self, diffusion_client: DiffusionClient):
+        self.diffusion_client = diffusion_client
+
+    def add_caption_box(self, image: Image.Image, text: str) -> Image.Image:
+        img_with_border = ImageOps.expand(image, border=15, fill="black")
+        img_copy = img_with_border.convert("RGBA")
+        overlay = Image.new("RGBA", img_copy.size, (255, 255, 255, 0))
+        draw = ImageDraw.Draw(overlay)
+
+        try:
+            font = ImageFont.truetype(FONT_PATH, 32)
+        except Exception:
+            font = ImageFont.load_default()
+
+        wrapped_text = textwrap.wrap(text, width=55)
+        line_height = 40
+        text_total_height = len(wrapped_text) * line_height
+
+        box_bottom = img_copy.size[1] - 30
+        box_top = box_bottom - text_total_height - 30
+        draw.rectangle(
+            [(30, box_top), (img_copy.size[0] - 30, box_bottom)],
+            fill=(255, 255, 255, 235),
+            outline=(0, 0, 0, 255),
+            width=6,
+        )
+
+        y_text = box_top + 15
+        for line in wrapped_text:
+            draw.text((45, y_text), line, font=font, fill="black")
+            y_text += line_height
+
+        return Image.alpha_composite(img_copy, overlay).convert("RGB")
+
+    def generate_image(
+        self,
+        prompt: str,
+        paragraph: str,
+        num_inference_steps: int = 5,
+        guidance_scale: float = 2.0,
+        size: int = 1024,
+    ) -> Image.Image | None:
+        raw_img = self.diffusion_client.gen_image(
+            prompt=prompt,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            width=size,
+            height=size,
+        )
+        if raw_img is not None:
+            return self.add_caption_box(raw_img, paragraph)
+        return None

src/image_prompt_generator.py

@@ -0,0 +1,50 @@
+import re
+import logging
+from src.llm import LLMClient
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class ImagePromptGenerator:
+    def __init__(self, llm_client: LLMClient):
+        self.llm_client = llm_client
+
+    def parse_paragraphs(self, raw_text: str) -> list[str]:
+        paragraphs = []
+        pattern = r"<\|paragraph\|>(.*?)(?=<\|paragraph\|>|<\|eos\|>|$)"
+        for match in re.finditer(pattern, raw_text, re.DOTALL):
+            content = match.group(1).strip()
+            if content:
+                paragraphs.append(content)
+        return paragraphs
+
+    def create_visual_prompt(
+        self,
+        paragraph: str,
+        full_story: str,
+        style: str = (
+            "whimsical fable book illustration, highly detailed, "
+            "vibrant colors, fantasy art style"
+        ),
+    ) -> str:
+        if self.llm_client is None:
+            logger.warning("LLM Client does not exist. Return fallback prompt.")
+            return f"{paragraph}, {style}"
+
+        translation_prompt = f"""You are an expert art director. Given the full context of a Vietnamese story, translate the specific paragraph into a highly short descriptive English image prompt.
+Maintain visual consistency of characters, objects, and environments based on the overall story.
+Focus only on the visual elements (characters, actions, environment) of the specific paragraph.
+Do NOT include any explanations or conversational text. Return ONLY the English translation, under 50 words.
+
+Full Story Context: "{full_story}"
+
+Specific Paragraph to visualize: "{paragraph}"
+
+English Image Prompt:"""
+
+        english_translation = self.llm_client.generate(
+            translation_prompt, max_new_tokens=100
+        )
+        english_translation = english_translation.strip(' "\'\n')
+        return f"{english_translation}, {style}"

src/llm.py

@@ -0,0 +1,93 @@
+import torch
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    BitsAndBytesConfig,
+    pipeline,
+)
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class LLMClient:
+    def __init__(
+        self,
+        model_id: str = "meta-llama/Meta-Llama-3-8B-Instruct",
+        hf_token: str | None = None,
+        cache_dir: str = "./cache",
+    ):
+        self.model_id = model_id
+        self.hf_token = hf_token
+        self.cache_dir = cache_dir
+        self.model = None
+        self.generator = None
+        self.tokenizer = None
+
+    def load_model(self):
+        if self.model is not None:
+            logger.info("LLM %s already loaded. Skipping.", self.model_id)
+            return
+
+        logger.info("Loading LLM %s...", self.model_id)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.model_id, token=self.hf_token, cache_dir=self.cache_dir
+        )
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        llama3_template = (
+            "{% set loop_messages = messages %}"
+            "{% for message in loop_messages %}"
+            "{% if message['role'] == 'system' %}"
+            "{{ '<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n' "
+            "+ message['content'] + '<|eot_id|>' }}"
+            "{% elif message['role'] == 'user' %}"
+            "{{ '<|start_header_id|>user<|end_header_id|>\n\n' "
+            "+ message['content'] + '<|eot_id|>' }}"
+            "{% elif message['role'] == 'assistant' %}"
+            "{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' "
+            "+ message['content'] + '<|eot_id|>' }}"
+            "{% endif %}{% endfor %}"
+            "{% if add_generation_prompt %}"
+            "{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}"
+            "{% endif %}"
+        )
+        self.tokenizer.chat_template = llama3_template
+
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.model_id,
+            device_map="auto",
+            quantization_config=bnb_config,
+            token=self.hf_token,
+            cache_dir=self.cache_dir,
+        )
+        self.generator = pipeline(
+            "text-generation", model=self.model, tokenizer=self.tokenizer
+        )
+        logger.info("LLM loaded successfully.")
+
+    def generate(self, prompt: str, max_new_tokens: int = 512) -> str:
+        if self.generator is None:
+            self.load_model()
+
+        messages = [
+            {"role": "system", "content": "You are a helpful and precise assistant."},
+            {"role": "user", "content": prompt},
+        ]
+        prompt_text = self.tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        outputs = self.generator(
+            prompt_text,
+            max_new_tokens=max_new_tokens,
+            return_full_text=False,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=self.tokenizer.eos_token_id,
+        )
+        return outputs[0]["generated_text"].strip()

src/stable_diffusion.py

@@ -0,0 +1,56 @@
+import torch
+import logging
+from diffusers import StableDiffusionXLPipeline
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class DiffusionClient:
+    def __init__(
+        self,
+        huggingface_path: str = "Lykon/dreamshaper-xl-v2-turbo",
+        hf_token: str | None = None,
+        cache_dir: str | None = "./cache",
+    ):
+        self.huggingface_path = huggingface_path
+        self.hf_token = hf_token
+        self.cache_dir = cache_dir
+        self.pipeline = None
+
+    def load_model(self):
+        logger.info("Loading SDXL model...")
+        self.pipeline = StableDiffusionXLPipeline.from_pretrained(
+            self.huggingface_path,
+            torch_dtype=torch.float16,
+            variant="fp16",
+            use_safetensors=True,
+            token=self.hf_token,
+            cache_dir=self.cache_dir,
+        )
+        self.pipeline.enable_model_cpu_offload()
+        self.pipeline.enable_vae_slicing()
+        self.pipeline.enable_vae_tiling()
+        logger.info("Loaded SDXL model %s successfully.", self.huggingface_path)
+
+    def gen_image(
+        self,
+        prompt: str,
+        negative_prompt: str = "",
+        num_inference_steps: int = 5,
+        guidance_scale: float = 2.0,
+        width: int = 1024,
+        height: int = 1024,
+    ):
+        if self.pipeline is None:
+            self.load_model()
+
+        image = self.pipeline(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            width=width,
+            height=height,
+        ).images[0]
+        return image

src/story_generator.py

@@ -0,0 +1,32 @@
+import logging
+from src.llm import LLMClient
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class StoryGenerator:
+    def __init__(self, llm_client: LLMClient):
+        self.model = llm_client
+
+    def load(self):
+        logger.info("Initializing Story Generator...")
+        self.model.load_model()
+
+    def create_story_prompt(self, user_story: str) -> str:
+        return f"""Hãy mở rộng đoạn tóm tắt dưới đây thành một câu chuyện ngụ ngôn.
+Tóm tắt câu chuyện:
+{user_story}
+
+Bạn phải định dạng đầu ra chính xác như sau, nội dung từng đoạn phải bắt đầu bằng một thẻ <|paragraph|>:
+<|paragraph|> [Nội dung đoạn 1] <|paragraph|> [Nội dung đoạn 2]
+
+Chú ý:
+- Không sinh thêm thẻ nào ngoài các thẻ <|paragraph|> đã được quy định.
+- Phong cách kể chuyện ngụ ngôn, không cần sinh tên cụ thể cho nhân vật.
+- Nội dung sinh ra là Tiếng Việt, giới hạn dưới 20 từ mỗi đoạn.
+"""
+
+    def gen_story_structured(self, user_story: str, max_new_tokens: int = 400) -> str:
+        final_prompt = self.create_story_prompt(user_story)
+        return self.model.generate(final_prompt, max_new_tokens=max_new_tokens)

static/Roboto-Regular.ttf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d7598e12c5dbef095ff8272cfc55da0250bd07fbdecbac8a530b9b277872a134
+size 488584