Update app.py

app.py

@@ -18,21 +18,17 @@ DEFAULT_CUSTOM_WEIGHTS_PATH = "PosterCraft/PosterCraft-v1_RL"
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-# No need to manually set CUDA device on Spaces
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-
 logging.basicConfig(
     level=logging.INFO,
     format="%(asctime)s - %(levelname)s - %(message)s",
 )
 
 # ------------------------------------------------------------------
-# 2. Model Download Function (Referencing JarvisIR's approach)
+# 2. Model Download Function (CPU only)
 # ------------------------------------------------------------------
 def download_model_weights(target_dir, repo_id, subdir=None):
     """
-    Download model weights to specified directory
+    Download model weights to specified directory (CPU operation)
     
     Args:
         target_dir (str): Local target directory
@@ -51,7 +47,6 @@ def download_model_weights(target_dir, repo_id, subdir=None):
 
     try:
         if subdir:
-            # If subdirectory is specified, only download that subdirectory
             snapshot_download(
                 repo_id=repo_id,
                 repo_type="model",
@@ -61,7 +56,6 @@ def download_model_weights(target_dir, repo_id, subdir=None):
             )
             src_dir = os.path.join(tmp_dir, subdir)
         else:
-            # Download entire repository
             snapshot_download(
                 repo_id=repo_id,
                 repo_type="model",
@@ -70,7 +64,6 @@ def download_model_weights(target_dir, repo_id, subdir=None):
             )
             src_dir = tmp_dir
 
-        # Copy to target directory
         if os.path.exists(src_dir):
             shutil.copytree(src_dir, target_dir)
             logging.info(f"Successfully downloaded {repo_id} to {target_dir}")
@@ -80,22 +73,50 @@ def download_model_weights(target_dir, repo_id, subdir=None):
     except Exception as e:
         logging.error(f"Download failed: {e}")
     finally:
-        # Clean up temporary directory
         if os.path.exists(tmp_dir):
             shutil.rmtree(tmp_dir)
 
 # ------------------------------------------------------------------
-# 3. Qwen Prompt Rewriting Agent
+# 3. Pre-download models (CPU operation)
 # ------------------------------------------------------------------
-class QwenRecapAgent:
-    def __init__(self, model_path: str):
-        self.model_path = model_path
-        self.tokenizer = None
-        self.model = None
-        self.is_loaded = False
-        
-        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:
+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):
+        logging.info("Downloading custom Transformer weights...")
+        download_model_weights(custom_weights_local, DEFAULT_CUSTOM_WEIGHTS_PATH)
+    
+    # Download Qwen model
+    qwen_local = "local_weights/Qwen3-8B"
+    if not os.path.exists(qwen_local):
+        logging.info("Downloading Qwen model...")
+        download_model_weights(qwen_local, DEFAULT_QWEN_MODEL_PATH)
+    
+    logging.info("Model download check completed")
+
+# Pre-download models at startup (CPU)
+ensure_models_downloaded()
+
+# ------------------------------------------------------------------
+# 4. Qwen Prompt Rewriting Agent
+# ------------------------------------------------------------------
+def create_qwen_agent(model_path):
+    """Create Qwen agent inside GPU context"""
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        torch_dtype=torch.bfloat16,
+        device_map="auto"
+    )
+    return tokenizer, model
+
+def recap_prompt(tokenizer, model, text):
+    """Recap prompt using Qwen model"""
+    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)
@@ -103,12 +124,14 @@ Identify the key elements in the user's prompt. Do not miss any details.
 - **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.
@@ -121,142 +144,44 @@ Elaborate on each core requirement to create a rich description.
     - **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 _load_model(self):
-        """Lazy load model"""
-        if not self.is_loaded:
-            logging.info(f"Loading Qwen model: {self.model_path}")
-            
-            # Ensure model files exist, if not download from Hub
-            if not os.path.exists(self.model_path):
-                download_model_weights(self.model_path, DEFAULT_QWEN_MODEL_PATH)
-            
-            self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
-            self.model = AutoModelForCausalLM.from_pretrained(
-                self.model_path,
-                torch_dtype=torch_dtype,
-                device_map="auto"
-            )
-            self.is_loaded = True
-
-    def recap(self, text: str) -> str:
-        try:
-            self._load_model()
-            
-            messages = [
-                {"role": "user", "content": self.prompt_template.format(brief_description=text)}
-            ]
-            chat = self.tokenizer.apply_chat_template(
-                messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
-            )
-            inputs = self.tokenizer([chat], return_tensors="pt").to(self.model.device)
-            
-            with torch.no_grad():
-                ids = self.model.generate(
-                    **inputs, max_new_tokens=1024, temperature=0.6, do_sample=True
-                )
-            out = self.tokenizer.decode(
-                ids[0][len(inputs.input_ids[0]):], skip_special_tokens=True
-            ).strip()
-            
-            if "</think>" in out:
-                out = out.split("</think>")[-1].strip()
-            return out or text
-        except Exception as e:
-            logging.warning(f"Recap failed: {e}. Using original prompt.")
-            return text
-
-# ------------------------------------------------------------------
-# 4. Poster Generator
-# ------------------------------------------------------------------
-class PosterGenerator:
-    def __init__(self):
-        self.pipeline = None
-        self.qwen = None
-        self.is_loaded = False
-
-    def _load_models(self):
-        """Lazy load all models"""
-        if not self.is_loaded:
-            logging.info("Starting model loading...")
-            
-            # Download custom weights (if not exists)
-            custom_weights_local = "local_weights/PosterCraft-v1_RL"
-            if not os.path.exists(custom_weights_local):
-                logging.info("Downloading custom Transformer weights...")
-                download_model_weights(custom_weights_local, DEFAULT_CUSTOM_WEIGHTS_PATH)
-            
-            # Load FLUX pipeline
-            logging.info("Loading FLUX pipeline...")
-            self.pipeline = FluxPipeline.from_pretrained(
-                DEFAULT_PIPELINE_PATH, 
-                torch_dtype=torch_dtype
-            )
-            
-            # Load custom Transformer
-            if os.path.exists(custom_weights_local):
-                try:
-                    logging.info("Loading custom Transformer...")
-                    transformer = FluxTransformer2DModel.from_pretrained(
-                        custom_weights_local, 
-                        torch_dtype=torch_dtype
-                    )
-                    self.pipeline.transformer = transformer
-                    logging.info("Custom Transformer loaded successfully")
-                except Exception as e:
-                    logging.warning(f"Custom weights loading failed: {e}, using default weights")
-            
-            # Enable memory optimization
-            self.pipeline.enable_model_cpu_offload()
-            
-            # Initialize Qwen (lazy loading)
-            qwen_local = "local_weights/Qwen3-8B"
-            if not os.path.exists(qwen_local):
-                logging.info("Downloading Qwen model...")
-                download_model_weights(qwen_local, DEFAULT_QWEN_MODEL_PATH)
-            
-            self.qwen = QwenRecapAgent(qwen_local)
-            
-            self.is_loaded = True
-            logging.info("All models loaded successfully")
 
-    def generate(self, prompt, enable_recap, **kwargs):
-        """Generate poster with given parameters"""
-        final_prompt = prompt
-        if enable_recap:
-            if not self.qwen:
-                raise gr.Error("Recap is enabled, but the recap model is not available. Check model path.")
-            final_prompt = self.qwen.recap(prompt)
-
-        generator = torch.Generator(device="cpu").manual_seed(kwargs['seed'])
+    try:
+        messages = [
+            {"role": "user", "content": prompt_template.format(brief_description=text)}
+        ]
+        chat = tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
+        )
+        inputs = tokenizer([chat], return_tensors="pt").to(model.device)
         
-        with torch.inference_mode():
-            image = self.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]
+        with torch.no_grad():
+            ids = model.generate(
+                **inputs, max_new_tokens=1024, temperature=0.6, do_sample=True
+            )
+        out = tokenizer.decode(
+            ids[0][len(inputs.input_ids[0]):], skip_special_tokens=True
+        ).strip()
         
-        return image, final_prompt
-
-# Global instance
-poster_gen = PosterGenerator()
+        if "</think>" in out:
+            out = out.split("</think>")[-1].strip()
+        return out or text
+    except Exception as e:
+        logging.warning(f"Recap failed: {e}. Using original prompt.")
+        return text
 
 # ------------------------------------------------------------------
 # 5. ZeroGPU Inference Function
 # ------------------------------------------------------------------
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=300)  # 增加到5分钟，给模型加载更多时间
 def generate_image_interface(
     original_prompt, enable_recap, height, width, 
     num_inference_steps, guidance_scale, seed_input,
@@ -268,45 +193,94 @@ def generate_image_interface(
     if width > MAX_IMAGE_SIZE or height > MAX_IMAGE_SIZE:
         raise gr.Error(f"Maximum resolution limit is {MAX_IMAGE_SIZE}×{MAX_IMAGE_SIZE}")
 
-    progress(0, desc="Loading models...")
-    
     try:
         actual_seed = int(seed_input) if seed_input and seed_input > 0 else random.randint(1, 2**32 - 1)
         
-        # Ensure models are loaded
-        poster_gen._load_models()
+        progress(0.1, desc="Loading FLUX pipeline...")
         
-        image, final_prompt = poster_gen.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
+        # Load FLUX pipeline
+        pipeline = FluxPipeline.from_pretrained(
+            DEFAULT_PIPELINE_PATH, 
+            torch_dtype=torch.bfloat16
         )
         
+        progress(0.2, desc="Loading custom transformer...")
+        
+        # Load custom transformer if available
+        custom_weights_local = "local_weights/PosterCraft-v1_RL"
+        if os.path.exists(custom_weights_local):
+            try:
+                transformer = FluxTransformer2DModel.from_pretrained(
+                    custom_weights_local, 
+                    torch_dtype=torch.bfloat16
+                )
+                pipeline.transformer = transformer
+                logging.info("Custom Transformer loaded successfully")
+            except Exception as e:
+                logging.warning(f"Custom weights loading failed: {e}, using default weights")
+        
+        # Move pipeline to GPU
+        pipeline = pipeline.to("cuda")
+        
+        final_prompt = original_prompt
+        
+        if enable_recap:
+            progress(0.4, desc="Loading Qwen model for prompt enhancement...")
+            
+            qwen_local = "local_weights/Qwen3-8B"
+            if os.path.exists(qwen_local):
+                try:
+                    tokenizer, model = create_qwen_agent(qwen_local)
+                    final_prompt = recap_prompt(tokenizer, model, original_prompt)
+                    progress(0.6, desc="Prompt enhanced, starting generation...")
+                    
+                    # Clean up Qwen model to free memory
+                    del tokenizer, model
+                    torch.cuda.empty_cache()
+                except Exception as e:
+                    logging.warning(f"Qwen model failed: {e}, using original prompt")
+                    final_prompt = original_prompt
+            else:
+                logging.warning("Qwen model not found, using original prompt")
+                final_prompt = original_prompt
+        
+        progress(0.7, desc="Generating image...")
+        
+        # Generate image
+        generator = torch.Generator(device="cuda").manual_seed(actual_seed)
+        
+        with torch.inference_mode():
+            image = 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="Generation complete!")
+        
         status_log = f"Seed: {actual_seed} | Generation complete."
-        progress(1, desc="Generation complete!")
         return image, final_prompt, status_log
         
     except Exception as e:
         logging.error(f"Generation failed: {e}")
-        raise gr.Error(f"An error occurred: {e}")
+        raise gr.Error(f"An error occurred: {str(e)}")
 
 # ------------------------------------------------------------------
-# 6. Gradio Interface (Similar to demo format)
+# 6. Gradio Interface
 # ------------------------------------------------------------------
 with gr.Blocks(theme=gr.themes.Soft(), title="PosterCraft") as demo:
     gr.Markdown("# PosterCraft-v1.0")
-    gr.Markdown(f"Running on: **{device}** | Base Pipeline: **{DEFAULT_PIPELINE_PATH}**")
+    gr.Markdown(f"Base Pipeline: **{DEFAULT_PIPELINE_PATH}**")
     gr.Markdown("⚠️ **First use requires model download, please wait about 10-15 minutes**")
 
     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.")
+            enable_recap_checkbox = gr.Checkbox(label="Enable Prompt Recap", value=True, info="Uses Qwen3-8B for prompt enhancement")
             
             with gr.Row():
                 width_input = gr.Slider(label="Width", minimum=256, maximum=2048, value=832, step=64)
@@ -333,4 +307,4 @@ with gr.Blocks(theme=gr.themes.Soft(), title="PosterCraft") as demo:
     generate_button.click(fn=generate_image_interface, inputs=inputs_list, outputs=outputs_list)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()