app.py

# -*- coding: utf-8 -*-
import os
import random
import logging
import numpy as np
import gradio as gr
import spaces
import torch
from huggingface_hub import login, whoami

# ─────────────────────────────────────────────
# 全局配置
# ─────────────────────────────────────────────
HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")

DEFAULT_PIPELINE_PATH = "black-forest-labs/FLUX.1-dev"
DEFAULT_QWEN_MODEL_PATH = "Qwen/Qwen3-8B"
DEFAULT_CUSTOM_WEIGHTS_PATH = "PosterCraft/PosterCraft-v1_RL"

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

# 认证状态
auth_status = "🔴 Not Authenticated"
if HF_TOKEN:
    try:
        login(token=HF_TOKEN, add_to_git_credential=True)
        user_info = whoami(HF_TOKEN)
        auth_status = f"✅ Authenticated as {user_info['name']}"
        logging.info(f"Successfully authenticated with Hugging Face as {user_info['name']}")
    except Exception as e:
        logging.error(f"HF authentication failed: {e}")
        auth_status = f"🔴 Authentication Error: {str(e)}"

# ─────────────────────────────────────────────
# 在 GPU 子进程 import 阶段就把大模型读进显存
# 只在 GPU 进程执行，CPU 主进程跳过
# ─────────────────────────────────────────────

# 检测是否在 GPU 环境中运行
def is_gpu_available():
    """检测当前环境是否支持 GPU"""
    try:
        import torch
        return torch.cuda.is_available()
    except ImportError:
        return False

if is_gpu_available():
    from diffusers import FluxPipeline, FluxTransformer2DModel
    from transformers import AutoModelForCausalLM, AutoTokenizer
    
    print("⏱️  [GPU init] Loading FLUX pipeline...")
    FLUX_PIPELINE = FluxPipeline.from_pretrained(
        DEFAULT_PIPELINE_PATH,
        torch_dtype=torch.bfloat16,
        token=HF_TOKEN
    ).to("cuda")

    print("⏱️  [GPU init] Loading PosterCraft transformer...")
    POSTERCRAFT_TRANSFORMER = FluxTransformer2DModel.from_pretrained(
        DEFAULT_CUSTOM_WEIGHTS_PATH,
        torch_dtype=torch.bfloat16,
        token=HF_TOKEN
    ).to("cuda")
    FLUX_PIPELINE.transformer = POSTERCRAFT_TRANSFORMER

    print("⏱️  [GPU init] Loading Qwen model...")
    QWEN_TOKENIZER = AutoTokenizer.from_pretrained(
        DEFAULT_QWEN_MODEL_PATH, 
        token=HF_TOKEN, 
        trust_remote_code=True,
        use_fast=True
    )
    QWEN_MODEL = AutoModelForCausalLM.from_pretrained(
        DEFAULT_QWEN_MODEL_PATH,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        token=HF_TOKEN,
        trust_remote_code=True,
    )
    
    print("✅ [GPU init] All models loaded successfully!")

# ─────────────────────────────────────────────
# Qwen 提示词增强函数
# ─────────────────────────────────────────────
def enhance_prompt_with_qwen(original_prompt):
    """使用预加载的 Qwen 模型增强提示词"""
    if not is_gpu_available():
        return original_prompt
        
    prompt_template = """You are an expert poster prompt designer. Your task is to rewrite a user's short poster prompt into a detailed and vivid long-format prompt. Follow these steps carefully:

**Step 1: Analyze the Core Requirements**
Identify the key elements in the user's prompt. Do not miss any details.
- **Subject:** What is the main subject? (e.g., a person, an object, a scene)
- **Style:** What is the visual style? (e.g., photorealistic, cartoon, vintage, minimalist)
- **Text:** Is there any text, like a title or slogan?
- **Color Palette:** Are there specific colors mentioned?
- **Composition:** Are there any layout instructions?

**Step 2: Expand and Add Detail**
Elaborate on each core requirement to create a rich description.
- **Do Not Omit:** You must include every piece of information from the original prompt.
- **Enrich with Specifics:** Add professional and descriptive details.
    - **Example:** If the user says "a woman with a bow", you could describe her as "a young woman with a determined expression, holding a finely crafted wooden longbow, with an arrow nocked and ready to fire."
- **Fill in the Gaps:** If the original prompt is simple (e.g., "a poster for a coffee shop"), use your creativity to add fitting details. You might add "The poster features a top-down view of a steaming latte with delicate art on its foam, placed on a rustic wooden table next to a few scattered coffee beans."

**Step 3: Handle Text Precisely**
- **Identify All Text Elements:** Carefully look for any text mentioned in the prompt. This includes:
    - **Explicit Text:** Subtitles, slogans, or any text in quotes.
    - **Implicit Titles:** The name of an event, movie, or product is often the main title. For example, if the prompt is "generate a 'Inception' poster ...", the title is "Inception".
- **Rules for Text:**
    - **If Text Exists:**
        - You must use the exact text identified from the prompt.
        - Do NOT add new text or delete existing text.
        - Describe each text's appearance (font, style, color, position). Example: `The title 'Inception' is written in a bold, sans-serif font, integrated into the cityscape.`
    - **If No Text Exists:**
        - Do not add any text elements. The poster must be purely visual.
- Most posters have titles. When a title exists, you must extend the title's description. Only when you are absolutely sure that there is no text to render, you can allow the extended prompt not to render text.

**Step 4: Final Output Rules**
- **Output ONLY the rewritten prompt.** No introductions, no explanations, no "Here is the prompt:".
- **Use a descriptive and confident tone.** Write as if you are describing a finished, beautiful poster.
- **Keep it concise.** The final prompt should be under 300 words.

---
**User Prompt:**
{brief_description}"""

    try:
        full_prompt = prompt_template.format(brief_description=original_prompt)
        messages = [{"role": "user", "content": full_prompt}]
        
        text = QWEN_TOKENIZER.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
        model_inputs = QWEN_TOKENIZER([text], return_tensors="pt").to(QWEN_MODEL.device)
        
        with torch.no_grad():
            generated_ids = QWEN_MODEL.generate(
                **model_inputs,
                max_new_tokens=512,          # ← 从 4096 改成 512 够用
                temperature=0.6,
                top_p=0.9,
                do_sample=True,
                eos_token_id=QWEN_TOKENIZER.eos_token_id,
            )
        
        output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
        full_response = QWEN_TOKENIZER.decode(output_ids, skip_special_tokens=True)
        
        # 提取最终答案
        if "</think>" in full_response:
            final_answer = full_response.split("</think>")[-1].strip()
        elif "<think>" not in full_response:
            final_answer = full_response.strip()
        else:
            final_answer = original_prompt
            
        return final_answer if final_answer else original_prompt
        
    except Exception as e:
        logging.error(f"Qwen enhancement failed: {e}")
        return original_prompt

# ─────────────────────────────────────────────
# 主要生成函数（使用预加载的模型）
# ─────────────────────────────────────────────
@spaces.GPU(duration=300)
def generate_poster(
    original_prompt, 
    enable_recap, 
    height, 
    width, 
    num_inference_steps, 
    guidance_scale, 
    seed_input,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate poster using preloaded models"""
    if not original_prompt or not original_prompt.strip():
        return None, "❌ Prompt cannot be empty!", ""

    try:
        if not HF_TOKEN:
            return None, "❌ Error: HF_TOKEN not found, please configure authentication.", ""

        progress(0.1, desc="Starting generation...")

        # Determine final prompt
        final_prompt = original_prompt
        if enable_recap:
            progress(0.2, desc="Enhancing prompt...")
            final_prompt = enhance_prompt_with_qwen(original_prompt)
        
        # Determine seed
        actual_seed = int(seed_input) if seed_input and seed_input != -1 else random.randint(1, 2**32 - 1)
        
        progress(0.3, desc="Generating image...")
        
        # Use preloaded FLUX pipeline to generate image
        generator = torch.Generator("cuda").manual_seed(actual_seed)
        
        with torch.inference_mode():
            image = FLUX_PIPELINE(
                prompt=final_prompt,
                generator=generator,
                num_inference_steps=int(num_inference_steps),
                guidance_scale=float(guidance_scale),
                width=int(width),
                height=int(height)
            ).images[0]
        
        progress(1.0, desc="Complete!")
        status_log = f"✅ Generation complete! Seed: {actual_seed}"
        return image, status_log, final_prompt

    except Exception as e:
        logging.error(f"Generation failed: {e}")
        return None, f"❌ Generation failed: {str(e)}", ""

# ─────────────────────────────────────────────
# Gradio Interface
# ─────────────────────────────────────────────
def create_interface():
    """Create Gradio interface"""
    
    with gr.Blocks(
        theme=gr.themes.Soft(),
        css="""
        .preserve-aspect-ratio img {
            object-fit: contain !important;
            max-height: 300px !important;
            width: auto !important;
        }
        """
    ) as demo:
        gr.HTML("<div style='text-align: center; width: 100%; margin: 20px 0;'><h1 style='margin: 0; font-size: 2.5em;'>PosterCraft-v1_RL</h1></div>")

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 1. Configuration")
                prompt_input = gr.Textbox(label="Prompt", lines=3, placeholder="Enter your creative prompt...")
                enable_recap_checkbox = gr.Checkbox(label="Enable Prompt Recap", value=True, info=f"Uses {DEFAULT_QWEN_MODEL_PATH} for rewriting.")
                
                with gr.Row():
                    width_input = gr.Slider(label="Width", minimum=256, maximum=2048, value=832, step=64)
                    height_input = gr.Slider(label="Height", minimum=256, maximum=2048, value=1216, step=64)
                gr.Markdown("Tip: Recommended size is 832x1216 for best results.")
                
                num_inference_steps_input = gr.Slider(label="Inference Steps", minimum=1, maximum=100, value=28, step=1)
                guidance_scale_input = gr.Slider(label="Guidance Scale (CFG)", minimum=0.0, maximum=20.0, value=3.5, step=0.1)
                seed_number_input = gr.Number(label="Seed", value=-1, minimum=-1, step=1, info="Leave blank or set to -1 for a random seed.")
                generate_button = gr.Button("Generate Image", variant="primary")

            with gr.Column(scale=1):
                gr.Markdown("### 2. Results")
                image_output = gr.Image(label="Generated Image", type="pil", show_download_button=True, height=340, container=False, elem_classes=["preserve-aspect-ratio"])
                recapped_prompt_output = gr.Textbox(label="Final Prompt Used", lines=5, interactive=False)
                status_output = gr.Textbox(label="Status Log", lines=4, interactive=False)

        inputs_list = [
            prompt_input, enable_recap_checkbox, height_input, width_input,
            num_inference_steps_input, guidance_scale_input, seed_number_input
        ]
        outputs_list = [image_output, recapped_prompt_output, status_output]
        
        generate_button.click(fn=generate_poster, inputs=inputs_list, outputs=outputs_list)
    
    return demo

# ─────────────────────────────────────────────
# 启动应用
# ─────────────────────────────────────────────
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_api=False
    )