Update app.py

app.py

@@ -5,21 +5,19 @@ 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")
+HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
 auth_status = "🔴 Not Authenticated"
 
-if hf_token:
+if HF_TOKEN:
     try:
-        login(token=hf_token, add_to_git_credential=True)
-        user_info = whoami(hf_token)
+        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:
@@ -64,8 +62,8 @@ def download_model_weights(target_dir, repo_id, subdir=None):
             "local_dir_use_symlinks": False,
         }
         
-        if hf_token:
-            download_kwargs["token"] = hf_token
+        if HF_TOKEN:
+            download_kwargs["token"] = HF_TOKEN
             
         if subdir:
             download_kwargs["allow_patterns"] = os.path.join(subdir, "**")
@@ -129,7 +127,7 @@ class QwenRecapAgent:
         try:
             self.tokenizer = AutoTokenizer.from_pretrained(
                 model_path, 
-                token=hf_token,
+                token=HF_TOKEN,
                 use_fast=True,  # 强制使用 fast tokenizer
                 trust_remote_code=True
             )
@@ -138,7 +136,7 @@ class QwenRecapAgent:
             logging.warning(f"Fast tokenizer failed, falling back to slow: {e}")
             self.tokenizer = AutoTokenizer.from_pretrained(
                 model_path, 
-                token=hf_token,
+                token=HF_TOKEN,
                 use_fast=False,
                 trust_remote_code=True
             )
@@ -148,8 +146,8 @@ class QwenRecapAgent:
             "device_map": device_map if device_map == "auto" else None,
             "trust_remote_code": True
         }
-        if hf_token:
-            model_kwargs["token"] = hf_token
+        if HF_TOKEN:
+            model_kwargs["token"] = HF_TOKEN
             
         self.model = AutoModelForCausalLM.from_pretrained(model_path, **model_kwargs)
         if device_map != "auto":
@@ -257,7 +255,7 @@ class PosterGenerator:
             self.pipeline = FluxPipeline.from_pretrained(
                 self.pipeline_path, 
                 torch_dtype=torch.bfloat16,
-                token=hf_token
+                token=HF_TOKEN
             )
             
             # 加载自定义权重
@@ -267,7 +265,7 @@ class PosterGenerator:
                 transformer = FluxTransformer2DModel.from_pretrained(
                     custom_weights_local, 
                     torch_dtype=torch.bfloat16,
-                    token=hf_token
+                    token=HF_TOKEN
                 )
                 self.pipeline.transformer = transformer
             elif self.custom_weights_path and os.path.exists(self.custom_weights_path):
@@ -275,7 +273,7 @@ class PosterGenerator:
                 transformer = FluxTransformer2DModel.from_pretrained(
                     self.custom_weights_path, 
                     torch_dtype=torch.bfloat16,
-                    token=hf_token
+                    token=HF_TOKEN
                 )
                 self.pipeline.transformer = transformer
             
@@ -319,7 +317,7 @@ def generate_image_interface(
         return None, "❌ Prompt cannot be empty!", ""
 
     try:
-        if not hf_token:
+        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")