Update app.py

app.py

@@ -1,3 +1,4 @@
+# -*- coding: utf-8 -*-
 import os
 import random
 import logging
@@ -7,25 +8,10 @@ import spaces
 import torch
 from huggingface_hub import login, whoami
 
-# ------------------------------------------------------------------
-# 1. Authentication and Global Configuration
-# ------------------------------------------------------------------
-# Authenticate with HF token
+# ─────────────────────────────────────────────
+# 全局配置
+# ─────────────────────────────────────────────
 HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_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"
@@ -34,126 +20,69 @@ 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)
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
+# 认证状态
+auth_status = "🔴 Not Authenticated"
+if HF_TOKEN:
     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")
-
+        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"Download failed: {e}")
-    finally:
-        if os.path.exists(tmp_dir):
-            shutil.rmtree(tmp_dir)
+        logging.error(f"HF authentication failed: {e}")
+        auth_status = f"🔴 Authentication Error: {str(e)}"
 
-# ------------------------------------------------------------------
-# 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}")
+# ─────────────────────────────────────────────
+# 在 GPU 子进程 import 阶段就把大模型读进显存
+# 只在 GPU 进程执行，CPU 主进程跳过
+# ─────────────────────────────────────────────
+if spaces.GPU.is_gpu():
+    from diffusers import FluxPipeline, FluxTransformer2DModel
+    from transformers import AutoModelForCausalLM, AutoTokenizer
     
-    # 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}")
+    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,
+    )
     
-    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
-
-        # 强制使用 Fast Tokenizer
-        try:
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                model_path, 
-                token=HF_TOKEN,
-                use_fast=True,  # 强制使用 fast tokenizer
-                trust_remote_code=True
-            )
-            logging.info("Successfully loaded fast tokenizer")
-        except Exception as e:
-            logging.warning(f"Fast tokenizer failed, falling back to slow: {e}")
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                model_path, 
-                token=HF_TOKEN,
-                use_fast=False,
-                trust_remote_code=True
-            )
+    print("✅ [GPU init] All models loaded successfully!")
+
+# ─────────────────────────────────────────────
+# Qwen 提示词增强函数
+# ─────────────────────────────────────────────
+def enhance_prompt_with_qwen(original_prompt):
+    """使用预加载的 Qwen 模型增强提示词"""
+    if not spaces.GPU.is_gpu():
+        return original_prompt
         
-        model_kwargs = {
-            "torch_dtype": torch.bfloat16, 
-            "device_map": device_map if device_map == "auto" else None,
-            "trust_remote_code": True
-        }
-        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:
+    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.
@@ -191,172 +120,95 @@ Elaborate on each core requirement to create a rich description.
 ---
 **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=4096, 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
+    try:
+        full_prompt = prompt_template.format(brief_description=original_prompt)
+        messages = [{"role": "user", "content": full_prompt}]
         
-        # 缓存变量
-        self.qwen_agent = None
-        self.pipeline = None
+        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)
         
-    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.no_grad():
+            generated_ids = QWEN_MODEL.generate(**model_inputs, max_new_tokens=4096, temperature=0.6)
         
-        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]
+        output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
+        full_response = QWEN_TOKENIZER.decode(output_ids, skip_special_tokens=True)
         
-        return image, final_prompt
+        # 提取最终答案
+        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
 
-# ------------------------------------------------------------------
-# 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,
+def generate_poster(
+    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!", ""
+        return None, "❌ 提示词不能为空！", ""
 
     try:
         if not HF_TOKEN:
-            return None, "❌ Error: HF_TOKEN not found. Please configure authentication.", ""
+            return None, "❌ 错误：未找到 HF_TOKEN，请配置认证。", ""
 
-        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
-            )
+        progress(0.1, desc="开始生成...")
+
+        # 确定最终提示词
+        final_prompt = original_prompt
+        if enable_recap:
+            progress(0.2, desc="增强提示词...")
+            final_prompt = enhance_prompt_with_qwen(original_prompt)
         
+        # 确定种子
         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...")
+        progress(0.3, desc="生成图像...")
         
-        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
-        )
+        # 使用预加载的 FLUX 管道生成图像
+        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]
         
-        status_log = f"✅ Generation complete! Seed: {actual_seed}"
+        progress(1.0, desc="完成！")
+        status_log = f"✅ 生成完成！种子：{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)}", ""
+        logging.error(f"生成失败：{e}")
+        return None, f"❌ 生成失败：{str(e)}", ""
 
-# ------------------------------------------------------------------
-# 7. Gradio Interface (保持你的原始风格)
-# ------------------------------------------------------------------
+# ─────────────────────────────────────────────
+# Gradio 界面
+# ─────────────────────────────────────────────
 def create_interface():
-    """Create Gradio interface"""
+    """创建 Gradio 界面"""
     
     with gr.Blocks(
         title="PosterCraft-v1.0",
@@ -371,50 +223,50 @@ def create_interface():
         <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
+                专业海报生成工具，基于 FLUX.1-dev 和定制微调权重
             </p>
         </div>
         """)
         
         with gr.Row():
-            gr.Markdown(f"**Base Pipeline:** `{DEFAULT_PIPELINE_PATH}`")
-            gr.Markdown(f"**Authentication Status:** {auth_status}")
+            gr.Markdown(f"**基础模型：** `{DEFAULT_PIPELINE_PATH}`")
+            gr.Markdown(f"**认证状态：** {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>
+            <p><strong>⚡ 首次生成需要加载模型（5-10分钟），后续生成会非常快！</strong></p>
         </div>
         """)
 
         with gr.Row():
             with gr.Column(scale=1):
-                gr.Markdown("### 1. Configuration")
+                gr.Markdown("### 1. 配置")
                 prompt_input = gr.Textbox(
-                    label="Poster Prompt", 
+                    label="海报提示词", 
                     lines=3, 
-                    placeholder="Enter your poster description...",
-                    value="A vintage travel poster for Paris, featuring the Eiffel Tower at sunset with warm golden lighting"
+                    placeholder="输入您的海报描述...",
+                    value="复古科幻电影海报，霓虹色彩和飞行汽车"
                 )
                 enable_recap_checkbox = gr.Checkbox(
-                    label="Enable Prompt Enhancement (Qwen3-8B)", 
+                    label="启用提示词增强 (Qwen3-8B)", 
                     value=True, 
-                    info="Use AI to enhance and expand your prompt"
+                    info="使用 AI 增强和扩展您的提示词"
                 )
                 
                 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)
+                    width_input = gr.Slider(label="宽度", minimum=256, maximum=MAX_IMAGE_SIZE, value=768, step=32)
+                    height_input = gr.Slider(label="高度", 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")
+                num_inference_steps_input = gr.Slider(label="推理步数", minimum=1, maximum=100, value=20, step=1)
+                guidance_scale_input = gr.Slider(label="引导强度", minimum=0.0, maximum=20.0, value=3.5, step=0.1)
+                seed_number_input = gr.Number(label="种子 (-1 随机)", value=-1, minimum=-1, step=1)
+                generate_button = gr.Button("🎨 生成海报", 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")
+                gr.Markdown("### 2. 结果")
+                image_output = gr.Image(label="生成的海报", type="pil", height=600)
+                status_output = gr.Textbox(label="生成状态", lines=2, interactive=False)
+                recapped_prompt_output = gr.Textbox(label="增强后的提示词", lines=5, interactive=False, info="用于生成的最终提示词")
 
         inputs_list = [
             prompt_input, enable_recap_checkbox, height_input, width_input,
@@ -422,28 +274,28 @@ def create_interface():
         ]
         outputs_list = [image_output, status_output, recapped_prompt_output]
         
-        generate_button.click(fn=generate_image_interface, inputs=inputs_list, outputs=outputs_list)
+        generate_button.click(fn=generate_poster, 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
-    ) 
+    )