app.py

import os
import random
import logging
import numpy as np
import gradio as gr
import spaces
import torch
from diffusers import FluxPipeline, FluxTransformer2DModel
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import login, whoami

# ------------------------------------------------------------------
# 1. Authentication and Global Configuration
# ------------------------------------------------------------------
# Authenticate with HF token
hf_token = os.getenv("HF_TOKEN")
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)}"
else:
    logging.warning("No HF_TOKEN found in environment variables")
    auth_status = "🔴 No HF_TOKEN found"

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

# ------------------------------------------------------------------
# 2. Model Download Function (CPU only)
# ------------------------------------------------------------------
def download_model_weights(target_dir, repo_id, subdir=None):
    """Download model weights to specified directory (CPU operation)"""
    from huggingface_hub import snapshot_download
    import shutil

    if os.path.exists(target_dir):
        logging.info(f"Directory {target_dir} already exists, skipping download")
        return

    tmp_dir = "hf_temp_download"
    os.makedirs(tmp_dir, exist_ok=True)

    try:
        download_kwargs = {
            "repo_id": repo_id,
            "repo_type": "model",
            "local_dir": tmp_dir,
            "local_dir_use_symlinks": False,
        }
        
        if hf_token:
            download_kwargs["token"] = hf_token
            
        if subdir:
            download_kwargs["allow_patterns"] = os.path.join(subdir, "**")
            
        snapshot_download(**download_kwargs)
        
        src_dir = os.path.join(tmp_dir, subdir) if subdir else tmp_dir

        if os.path.exists(src_dir):
            shutil.copytree(src_dir, target_dir)
            logging.info(f"Successfully downloaded {repo_id} to {target_dir}")
        else:
            logging.warning(f"Source directory {src_dir} does not exist")

    except Exception as e:
        logging.error(f"Download failed: {e}")
    finally:
        if os.path.exists(tmp_dir):
            shutil.rmtree(tmp_dir)

# ------------------------------------------------------------------
# 3. Pre-download models (CPU operation)
# ------------------------------------------------------------------
def ensure_models_downloaded():
    """Pre-download all models to avoid GPU timeout"""
    logging.info("Checking and downloading models if needed...")
    
    # Download custom weights
    custom_weights_local = "local_weights/PosterCraft-v1_RL"
    if not os.path.exists(custom_weights_local):
        try:
            logging.info("Downloading custom Transformer weights...")
            download_model_weights(custom_weights_local, DEFAULT_CUSTOM_WEIGHTS_PATH)
        except Exception as e:
            logging.warning(f"Failed to download custom weights: {e}")
    
    # Download Qwen model
    qwen_local = "local_weights/Qwen3-8B"
    if not os.path.exists(qwen_local):
        try:
            logging.info("Downloading Qwen model...")
            download_model_weights(qwen_local, DEFAULT_QWEN_MODEL_PATH)
        except Exception as e:
            logging.warning(f"Failed to download Qwen model: {e}")
    
    logging.info("Model download check completed")

# Pre-download models at startup (CPU)
ensure_models_downloaded()

# ------------------------------------------------------------------
# 4. Qwen Recap Agent (基于你的原始逻辑)
# ------------------------------------------------------------------
class QwenRecapAgent:
    def __init__(self, model_path, max_retries=3, retry_delay=2, device_map="auto"):
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.device = device_map

        self.tokenizer = AutoTokenizer.from_pretrained(model_path, token=hf_token)
        model_kwargs = {"torch_dtype": torch.bfloat16, "device_map": device_map if device_map == "auto" else None}
        if hf_token:
            model_kwargs["token"] = hf_token
        self.model = AutoModelForCausalLM.from_pretrained(model_path, **model_kwargs)
        if device_map != "auto":
             self.model.to(device_map)
        
        self.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}"""
    
    def recap_prompt(self, original_prompt):
        full_prompt = self.prompt_template.format(brief_description=original_prompt)
        messages = [{"role": "user", "content": full_prompt}]
        try:
            text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
            model_inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device)
            
            with torch.no_grad():
                generated_ids = self.model.generate(**model_inputs, max_new_tokens=1024, temperature=0.6)
            
            output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
            full_response = self.tokenizer.decode(output_ids, skip_special_tokens=True)
            final_answer = self._extract_final_answer(full_response)
            
            if final_answer:
                return final_answer.strip()
            
            logging.info("Qwen returned an empty answer. Using original prompt.")
            return original_prompt
        except Exception as e:
            logging.error(f"Qwen recap failed: {e}. Using original prompt.")
            return original_prompt

    def _extract_final_answer(self, full_response):
        if "</think>" in full_response:
            return full_response.split("</think>")[-1].strip()
        if "<think>" not in full_response:
            return full_response.strip()
        return None

# ------------------------------------------------------------------
# 5. Poster Generator Class (基于你的原始逻辑，但加上缓存)
# ------------------------------------------------------------------
class PosterGenerator:
    def __init__(self, pipeline_path, qwen_model_path, custom_weights_path, device):
        self.device = device
        self.pipeline_path = pipeline_path
        self.qwen_model_path = qwen_model_path
        self.custom_weights_path = custom_weights_path
        
        # 缓存变量
        self.qwen_agent = None
        self.pipeline = None
        
    def _load_qwen_agent(self):
        if self.qwen_agent is None:
            if not self.qwen_model_path:
                return None
            
            # 检查本地路径
            qwen_local = "local_weights/Qwen3-8B"
            model_path = qwen_local if os.path.exists(qwen_local) else self.qwen_model_path
            
            logging.info(f"Loading Qwen agent from {model_path}")
            self.qwen_agent = QwenRecapAgent(model_path=model_path, device_map=str(self.device))
        return self.qwen_agent

    def _load_flux_pipeline(self):
        if self.pipeline is None:
            logging.info("Loading FLUX pipeline...")
            self.pipeline = FluxPipeline.from_pretrained(
                self.pipeline_path, 
                torch_dtype=torch.bfloat16,
                token=hf_token
            )
            
            # 加载自定义权重
            custom_weights_local = "local_weights/PosterCraft-v1_RL"
            if os.path.exists(custom_weights_local):
                logging.info(f"Loading custom Transformer from directory: {custom_weights_local}")
                transformer = FluxTransformer2DModel.from_pretrained(
                    custom_weights_local, 
                    torch_dtype=torch.bfloat16,
                    token=hf_token
                )
                self.pipeline.transformer = transformer
            elif self.custom_weights_path and os.path.exists(self.custom_weights_path):
                logging.info(f"Loading custom Transformer from directory: {self.custom_weights_path}")
                transformer = FluxTransformer2DModel.from_pretrained(
                    self.custom_weights_path, 
                    torch_dtype=torch.bfloat16,
                    token=hf_token
                )
                self.pipeline.transformer = transformer
            
            self.pipeline.to(self.device)
        return self.pipeline

    def generate(self, prompt, enable_recap, **kwargs):
        final_prompt = prompt
        if enable_recap:
            qwen_agent = self._load_qwen_agent()
            if not qwen_agent:
                raise gr.Error("Recap is enabled, but the recap model is not available. Check model path.")
            final_prompt = qwen_agent.recap_prompt(prompt)

        pipeline = self._load_flux_pipeline()
        generator = torch.Generator(device=self.device).manual_seed(kwargs['seed'])
        
        with torch.inference_mode():
            image = pipeline(
                prompt=final_prompt,
                generator=generator,
                num_inference_steps=kwargs['num_inference_steps'],
                guidance_scale=kwargs['guidance_scale'],
                width=kwargs['width'],
                height=kwargs['height']
            ).images[0]
        
        return image, final_prompt

# ------------------------------------------------------------------
# 6. Main Generation Function (GPU) - 保持你的原始逻辑
# ------------------------------------------------------------------
@spaces.GPU(duration=300)
def generate_image_interface(
    original_prompt, enable_recap, height, width, 
    num_inference_steps, guidance_scale, seed_input,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate image using FLUX pipeline"""
    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.", ""

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # 全局生成器实例
        if not hasattr(generate_image_interface, 'generator'):
            generate_image_interface.generator = PosterGenerator(
                pipeline_path=DEFAULT_PIPELINE_PATH,
                qwen_model_path=DEFAULT_QWEN_MODEL_PATH,
                custom_weights_path=DEFAULT_CUSTOM_WEIGHTS_PATH,
                device=device
            )
        
        actual_seed = int(seed_input) if seed_input and seed_input != -1 else random.randint(1, 2**32 - 1)
        
        progress(0.1, desc="Starting generation...")
        
        image, final_prompt = generate_image_interface.generator.generate(
            prompt=original_prompt,
            enable_recap=enable_recap,
            height=int(height),
            width=int(width),
            num_inference_steps=int(num_inference_steps),
            guidance_scale=float(guidance_scale),
            seed=actual_seed
        )
        
        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)}", ""

# ------------------------------------------------------------------
# 7. Gradio Interface (保持你的原始风格)
# ------------------------------------------------------------------
def create_interface():
    """Create Gradio interface"""
    
    with gr.Blocks(
        title="PosterCraft-v1.0",
        theme=gr.themes.Soft(),
        css="""
        .main-container { max-width: 1200px; margin: 0 auto; }
        .status-box { padding: 10px; border-radius: 5px; margin: 10px 0; }
        """
    ) as demo:
        
        gr.HTML("""
        <div class="main-container">
            <h1 style="text-align: center; margin-bottom: 20px;">🎨 PosterCraft-v1.0</h1>
            <p style="text-align: center; color: #666; margin-bottom: 30px;">
                Professional poster generation with FLUX.1-dev and custom fine-tuned weights
            </p>
        </div>
        """)
        
        with gr.Row():
            gr.Markdown(f"**Base Pipeline:** `{DEFAULT_PIPELINE_PATH}`")
            gr.Markdown(f"**Authentication Status:** {auth_status}")
        
        gr.HTML("""
        <div class="status-box">
            <p><strong>⚠️ First generation requires model loading (5-10 minutes). Subsequent generations are much faster!</strong></p>
        </div>
        """)

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 1. Configuration")
                prompt_input = gr.Textbox(
                    label="Poster Prompt", 
                    lines=3, 
                    placeholder="Enter your poster description...",
                    value="A vintage travel poster for Paris, featuring the Eiffel Tower at sunset with warm golden lighting"
                )
                enable_recap_checkbox = gr.Checkbox(
                    label="Enable Prompt Enhancement (Qwen3-8B)", 
                    value=True, 
                    info="Use AI to enhance and expand your prompt"
                )
                
                with gr.Row():
                    width_input = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, value=768, step=32)
                    height_input = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, value=1024, step=32)
                
                num_inference_steps_input = gr.Slider(label="Inference Steps", minimum=1, maximum=100, value=20, step=1)
                guidance_scale_input = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=20.0, value=3.5, step=0.1)
                seed_number_input = gr.Number(label="Seed (-1 for random)", value=-1, minimum=-1, step=1)
                generate_button = gr.Button("🎨 Generate Poster", variant="primary", size="lg")

            with gr.Column(scale=1):
                gr.Markdown("### 2. Results")
                image_output = gr.Image(label="Generated Poster", type="pil", height=600)
                status_output = gr.Textbox(label="Generation Status", lines=2, interactive=False)
                recapped_prompt_output = gr.Textbox(label="Enhanced Prompt", lines=5, interactive=False, info="The final prompt used for generation")

        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, status_output, recapped_prompt_output]
        
        generate_button.click(fn=generate_image_interface, inputs=inputs_list, outputs=outputs_list)
        
        # Examples
        gr.Examples(
            examples=[
                ["A retro sci-fi movie poster with neon colors and flying cars"],
                ["An elegant art deco poster for a luxury hotel"],
                ["A minimalist concert poster with bold typography"],
                ["A vintage advertisement for organic coffee"],
            ],
            inputs=[prompt_input]
        )
    
    return demo

# ------------------------------------------------------------------
# 8. Launch Application
# ------------------------------------------------------------------
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_api=False
    )