{commit_message}

.official_space.py

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:51292da76658340750a198cba125f12668ff88f79248ffc920e168b645698b43
-size 24588

app.py

@@ -59,9 +59,6 @@ pipe.to("cuda", torch.bfloat16)
 
 print("Model loaded successfully!")
 
-# Initialize LLM for prompt enhancement
-llm_client = InferenceClient()
-
 # Vision-Language model for prompt enhancement
 VL_MODEL = "Qwen/Qwen3-VL-30B-A3B-Instruct"
 
@@ -93,9 +90,16 @@ def image_to_base64(image) -> str:
     image.save(buffered, format="JPEG", quality=85)
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
-def enhance_prompt(prompt: str, reference_image=None) -> str:
+def enhance_prompt(prompt: str, reference_image=None, oauth_token: str = None) -> str:
     """Enhance the prompt using a VL model, optionally with a reference image."""
+    if not oauth_token:
+        print("[Prompt Enhancement] No auth token provided")
+        return prompt
+    
     try:
+        # Create client with user's token
+        client = InferenceClient(token=oauth_token)
+        
         # Build user content based on whether image is provided
         if reference_image is not None:
             # Convert image to base64 for the API
@@ -118,7 +122,7 @@ def enhance_prompt(prompt: str, reference_image=None) -> str:
             {"role": "user", "content": user_content}
         ]
         
-        response = llm_client.chat_completion(
+        response = client.chat_completion(
             messages=messages,
             model=VL_MODEL,
             max_tokens=250,
@@ -130,7 +134,7 @@ def enhance_prompt(prompt: str, reference_image=None) -> str:
         # Remove any thinking tags if present
         if "<think>" in enhanced:
             enhanced = re.sub(r'<think>.*?</think>', '', enhanced, flags=re.DOTALL).strip()
-        print(f"[Prompt Enhancement] Model: {model}")
+        print(f"[Prompt Enhancement] Model: {VL_MODEL}")
         print(f"[Prompt Enhancement] Original: {prompt}")
         print(f"[Prompt Enhancement] Enhanced: {enhanced}")
         return enhanced
@@ -151,11 +155,13 @@ def infer(
     num_inference_steps,
     use_prompt_enhancement,
     reference_image,
+    oauth_token: gr.OAuthToken | None,
     progress=gr.Progress(track_tqdm=True),
 ):
     # Enhance prompt if requested
     if use_prompt_enhancement:
-        prompt = enhance_prompt(prompt, reference_image)
+        token = oauth_token.token if oauth_token else None
+        prompt = enhance_prompt(prompt, reference_image, token)
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -286,8 +292,9 @@ footer { display: none !important; }
 with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown("# Z-Image", elem_classes="title")
-
-        gr.Markdown("", elem_classes="spacer")
+        
+        # Login button for HF authentication
+        login_btn = gr.LoginButton(value="Sign in with Hugging Face")
         
         # Prompt
         prompt = gr.Textbox(