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)}"


def is_gpu_available():
    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!")

def enhance_prompt_with_qwen(original_prompt):
    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,          
                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=30)
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="Re-writing 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)}", ""

def create_interface():
    """Create Gradio interface"""
    
    custom_css = """
.gradio-container {
    background: linear-gradient(135deg, #3b4371 0%, #2d1b69 25%, #673ab7 50%, #8e24aa 75%, #6a1b9a 100%);
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
    min-height: 100vh;
}

.contain {
    background: rgba(255, 255, 255, 0.95);
    border-radius: 15px;
    padding: 25px;
    margin: 15px;
    box-shadow: 0 10px 30px rgba(0, 0, 0, 0.2);
    backdrop-filter: blur(10px);
}

.title-container {
    text-align: center;
    margin-bottom: 25px;
    padding: 20px;
    background: linear-gradient(135deg, #673ab7, #8e24aa);
    border-radius: 12px;
    box-shadow: 0 5px 20px rgba(103, 58, 183, 0.4);
}

.title-container h1 {
    color: white;
    font-size: 2.2em;
    font-weight: bold;
    margin: 0;
    text-shadow: 1px 1px 3px rgba(0, 0, 0, 0.3);
}

.info-bar {
    background: linear-gradient(135deg, #7c4dff, #6a1b9a);
    padding: 12px;
    border-radius: 8px;
    margin-bottom: 20px;
    color: white;
    text-align: center;
    font-weight: 500;
    box-shadow: 0 3px 12px rgba(124, 77, 255, 0.3);
}

.section-header {
    background: linear-gradient(135deg, #e1bee7, #d1c4e9);
    padding: 12px;
    border-radius: 8px;
    margin-bottom: 15px;
    border-left: 4px solid #673ab7;
}

.section-header h3 {
    margin: 0;
    color: #333;
    font-weight: 600;
}

.input-group {
    background: rgba(255, 255, 255, 0.85);
    padding: 18px;
    border-radius: 12px;
    margin-bottom: 15px;
    border: 1px solid rgba(103, 58, 183, 0.2);
    box-shadow: 0 3px 12px rgba(103, 58, 183, 0.1);
}

.result-section {
    background: rgba(255, 255, 255, 0.9);
    padding: 18px;
    border-radius: 12px;
    border: 1px solid rgba(103, 58, 183, 0.2);
    box-shadow: 0 3px 12px rgba(103, 58, 183, 0.1);
}

.tip-box {
    background: linear-gradient(135deg, #f3e5f5, #e8eaf6);
    padding: 10px;
    border-radius: 6px;
    margin: 8px 0;
    border-left: 3px solid #673ab7;
    color: #4a148c;
    font-weight: 500;
}

button.primary {
    background: linear-gradient(135deg, #673ab7, #8e24aa) !important;
    border: none !important;
    border-radius: 20px !important;
    padding: 12px 25px !important;
    color: white !important;
    font-weight: bold !important;
    font-size: 15px !important;
    box-shadow: 0 5px 15px rgba(103, 58, 183, 0.4) !important;
}

button.primary:hover {
    box-shadow: 0 8px 25px rgba(103, 58, 183, 0.6) !important;
    opacity: 0.9 !important;
    transform: translateY(-2px) !important;
}


label {
    color: #4a148c !important;
    font-weight: 600 !important;
}

input, textarea, select {
    border: 1px solid rgba(103, 58, 183, 0.3) !important;
    border-radius: 6px !important;
}

input:focus, textarea:focus, select:focus {
    border-color: #673ab7 !important;
    box-shadow: 0 0 0 2px rgba(103, 58, 183, 0.2) !important;
}

.gr-slider input[type="range"] {
    accent-color: #673ab7 !important;
}

input[type="checkbox"] {
    accent-color: #673ab7 !important;
}

.preserve-aspect-ratio img {
    object-fit: contain !important;
    width: auto !important;
    max-height: 512px !important;
}
"""
    
    with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
        with gr.Column(elem_classes="contain"):
            gr.HTML('<div class="title-container"><h1>🎨 PosterCraft-v1.0</h1></div>')

            with gr.Row():
                with gr.Column(scale=1, elem_classes="input-group"):
                    gr.HTML('<div class="section-header"><h3>⚙️ 1. Configuration</h3></div>')
                    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 Qwen3 for rewriting.")
                    gr.Examples(
                        examples=[
                            ["Urban Canvas Street Art Expo poster with bold graffiti-style lettering and dynamic colorful splashes"],
                            ["This poster for 'PixelPlay Retro Game Console' features the console with classic 8-bit game graphics, evoking nostalgia and fun with a vibrant, playful, and retro-gaming aesthetic."],
                            ["Poster about Mars Tourism Campaign, text:\"NEXT STOP MARS\\nBOOK YOUR TICKET NOW\", astronaut_on_red_planet, rocket_launch, sunrise_horizon_glow, retro_futurism_style, dust_clouds, panoramic_view, bold_headline_text, sci-fi_palette, highres, 16x9_ratio"],
                            ["This intriguing poster for \"CODE OF THE SAMURAI\" presents a stark contrast. On one side, a traditional samurai warrior in full armor, holding a katana, is depicted in a sepia-toned, historical style. On the other side, a futuristic cyborg warrior with glowing blue optics and sleek armor is shown in a cool, modern, digital style. The two figures are back-to-back, divided by a shimmering energy line. The title \"CODE OF THE SAMURAI\" is written in a font that blends traditional Japanese calligraphy with modern digital elements, in a metallic silver, positioned horizontally across the center where the two styles meet. The tagline, \"HONOR IS TIMELESS,\" is in a smaller, clean white sans-serif font at the bottom. The layout highlights the duality and the clash or merging of ancient traditions with future technology."]
                        ],
                        inputs=[prompt_input],
                        label="📝 Example Prompts",
                        examples_per_page=5
                    )
                    
                    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.HTML('<div class="tip-box">💡 Tip: Recommended size is 832x1216 for best results.</div>')
                    
                    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, elem_classes="result-section"):
                    gr.HTML('<div class="section-header"><h3>🎨 2. Results</h3></div>')
                    image_output = gr.Image(label="🖼️ Generated Image", type="pil", show_download_button=True, height=512, 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
    )