Update app.py

app.py

@@ -7,10 +7,23 @@ import spaces
 import torch
 from diffusers import FluxPipeline, FluxTransformer2DModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import login
 
 # ------------------------------------------------------------------
-# 1. Global Configuration
+# 1. Authentication and Global Configuration
 # ------------------------------------------------------------------
+# Authenticate with HF token
+hf_token = os.getenv("HF_TOKEN")
+if hf_token:
+    try:
+        login(token=hf_token, add_to_git_credential=True)
+        logging.info("Successfully authenticated with Hugging Face")
+    except Exception as e:
+        logging.error(f"HF authentication failed: {e}")
+        raise Exception("Authentication failed. Please check your HF_TOKEN.")
+else:
+    logging.warning("No HF_TOKEN found in environment variables")
+
 DEFAULT_PIPELINE_PATH = "black-forest-labs/FLUX.1-dev"
 DEFAULT_QWEN_MODEL_PATH = "Qwen/Qwen3-8B"
 DEFAULT_CUSTOM_WEIGHTS_PATH = "PosterCraft/PosterCraft-v1_RL"
@@ -46,23 +59,23 @@ def download_model_weights(target_dir, repo_id, subdir=None):
     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,
+        }
+        
+        # Add token if available
+        if hf_token:
+            download_kwargs["token"] = hf_token
+            
         if subdir:
-            snapshot_download(
-                repo_id=repo_id,
-                repo_type="model",
-                local_dir=tmp_dir,
-                allow_patterns=os.path.join(subdir, "**"),
-                local_dir_use_symlinks=False,
-            )
-            src_dir = os.path.join(tmp_dir, subdir)
-        else:
-            snapshot_download(
-                repo_id=repo_id,
-                repo_type="model",
-                local_dir=tmp_dir,
-                local_dir_use_symlinks=False,
-            )
-            src_dir = tmp_dir
+            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)
@@ -86,14 +99,20 @@ def ensure_models_downloaded():
     # 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)
+        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):
-        logging.info("Downloading Qwen model...")
-        download_model_weights(qwen_local, DEFAULT_QWEN_MODEL_PATH)
+        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")
 
@@ -105,12 +124,17 @@ ensure_models_downloaded()
 # ------------------------------------------------------------------
 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"
-    )
+    load_kwargs = {
+        "torch_dtype": torch.bfloat16,
+        "device_map": "auto"
+    }
+    
+    # Add token if available
+    if hf_token:
+        load_kwargs["token"] = hf_token
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_path, **load_kwargs)
+    model = AutoModelForCausalLM.from_pretrained(model_path, **load_kwargs)
     return tokenizer, model
 
 def recap_prompt(tokenizer, model, text):
@@ -181,7 +205,7 @@ Elaborate on each core requirement to create a rich description.
 # ------------------------------------------------------------------
 # 5. ZeroGPU Inference Function
 # ------------------------------------------------------------------
-@spaces.GPU(duration=300)  # 增加到5分钟，给模型加载更多时间
+@spaces.GPU(duration=300)
 def generate_image_interface(
     original_prompt, enable_recap, height, width, 
     num_inference_steps, guidance_scale, seed_input,
@@ -198,11 +222,14 @@ def generate_image_interface(
         
         progress(0.1, desc="Loading FLUX pipeline...")
         
-        # Load FLUX pipeline
-        pipeline = FluxPipeline.from_pretrained(
-            DEFAULT_PIPELINE_PATH, 
-            torch_dtype=torch.bfloat16
-        )
+        # Load FLUX pipeline with explicit token
+        load_kwargs = {
+            "torch_dtype": torch.bfloat16
+        }
+        if hf_token:
+            load_kwargs["token"] = hf_token
+            
+        pipeline = FluxPipeline.from_pretrained(DEFAULT_PIPELINE_PATH, **load_kwargs)
         
         progress(0.2, desc="Loading custom transformer...")
         
@@ -210,9 +237,12 @@ def generate_image_interface(
         custom_weights_local = "local_weights/PosterCraft-v1_RL"
         if os.path.exists(custom_weights_local):
             try:
+                transformer_kwargs = {"torch_dtype": torch.bfloat16}
+                if hf_token:
+                    transformer_kwargs["token"] = hf_token
+                    
                 transformer = FluxTransformer2DModel.from_pretrained(
-                    custom_weights_local, 
-                    torch_dtype=torch.bfloat16
+                    custom_weights_local, **transformer_kwargs
                 )
                 pipeline.transformer = transformer
                 logging.info("Custom Transformer loaded successfully")
@@ -274,6 +304,11 @@ def generate_image_interface(
 with gr.Blocks(theme=gr.themes.Soft(), title="PosterCraft") as demo:
     gr.Markdown("# PosterCraft-v1.0")
     gr.Markdown(f"Base Pipeline: **{DEFAULT_PIPELINE_PATH}**")
+    
+    # Show authentication status
+    auth_status = "🟢 Authenticated" if hf_token else "🔴 Not Authenticated"
+    gr.Markdown(f"Authentication Status: {auth_status}")
+    
     gr.Markdown("⚠️ **First use requires model download, please wait about 10-15 minutes**")
 
     with gr.Row():