Update app.py

app.py

@@ -7,22 +7,27 @@ import spaces
 import torch
 from diffusers import FluxPipeline, FluxTransformer2DModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from huggingface_hub import login
+from huggingface_hub import login, whoami
 
 # ------------------------------------------------------------------
 # 1. Authentication and Global Configuration
 # ------------------------------------------------------------------
 # Authenticate with HF token
 hf_token = os.getenv("HF_TOKEN")
+auth_status = "🔴 Not Authenticated"
+
 if hf_token:
     try:
         login(token=hf_token, add_to_git_credential=True)
-        logging.info("Successfully authenticated with Hugging Face")
+        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}")
-        raise Exception("Authentication failed. Please check your HF_TOKEN.")
+        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"
@@ -199,11 +204,11 @@ Elaborate on each core requirement to create a rich description.
             out = out.split("</think>")[-1].strip()
         return out or text
     except Exception as e:
-        logging.warning(f"Recap failed: {e}. Using original prompt.")
+        logging.error(f"Prompt recap failed: {e}")
         return text
 
 # ------------------------------------------------------------------
-# 5. ZeroGPU Inference Function
+# 5. Main Generation Function (GPU)
 # ------------------------------------------------------------------
 @spaces.GPU(duration=300)
 def generate_image_interface(
@@ -211,135 +216,246 @@ def generate_image_interface(
     num_inference_steps, guidance_scale, seed_input,
     progress=gr.Progress(track_tqdm=True),
 ):
-    if not original_prompt or not original_prompt.strip():
-        raise gr.Error("Prompt cannot be empty!")
+    """Generate image using FLUX pipeline"""
+    try:
+        # If no token available, return error message
+        if not hf_token:
+            return None, "❌ Error: HF_TOKEN not found. Please configure authentication.", ""
 
-    if width > MAX_IMAGE_SIZE or height > MAX_IMAGE_SIZE:
-        raise gr.Error(f"Maximum resolution limit is {MAX_IMAGE_SIZE}×{MAX_IMAGE_SIZE}")
+        # Set device and dtype
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        torch_dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
 
-    try:
-        actual_seed = int(seed_input) if seed_input and seed_input > 0 else random.randint(1, 2**32 - 1)
-        
+        # Initialize FLUX pipeline
         progress(0.1, desc="Loading FLUX pipeline...")
+        pipeline = FluxPipeline.from_pretrained(
+            DEFAULT_PIPELINE_PATH,
+            torch_dtype=torch_dtype,
+            device_map="balanced" if device.type == "cuda" else None,
+            token=hf_token
+        )
         
-        # 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...")
-        
-        # Load custom transformer if available
+        # Load custom transformer weights if available
         custom_weights_local = "local_weights/PosterCraft-v1_RL"
         if os.path.exists(custom_weights_local):
+            progress(0.3, desc="Loading custom transformer weights...")
             try:
-                transformer_kwargs = {"torch_dtype": torch.bfloat16}
-                if hf_token:
-                    transformer_kwargs["token"] = hf_token
-                    
-                transformer = FluxTransformer2DModel.from_pretrained(
-                    custom_weights_local, **transformer_kwargs
+                custom_transformer = FluxTransformer2DModel.from_pretrained(
+                    custom_weights_local,
+                    torch_dtype=torch_dtype,
+                    device_map="balanced" if device.type == "cuda" else None,
+                    token=hf_token
                 )
-                pipeline.transformer = transformer
-                logging.info("Custom Transformer loaded successfully")
+                pipeline.transformer = custom_transformer
+                logging.info("Custom transformer weights 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")
+                logging.warning(f"Failed to load custom transformer weights: {e}")
         
+        # Process prompt
         final_prompt = original_prompt
-        
         if enable_recap:
-            progress(0.4, desc="Loading Qwen model for prompt enhancement...")
-            
+            progress(0.5, desc="Processing prompt with Qwen...")
             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...")
-                    
+                    logging.info(f"Enhanced prompt: {final_prompt}")
                     # 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")
+                    logging.warning(f"Qwen processing failed: {e}")
                     final_prompt = original_prompt
             else:
-                logging.warning("Qwen model not found, using original prompt")
-                final_prompt = original_prompt
-        
-        progress(0.7, desc="Generating image...")
+                # Fallback to online Qwen model
+                try:
+                    tokenizer, model = create_qwen_agent(DEFAULT_QWEN_MODEL_PATH)
+                    final_prompt = recap_prompt(tokenizer, model, original_prompt)
+                    del tokenizer, model
+                    torch.cuda.empty_cache()
+                except Exception as e:
+                    logging.warning(f"Online Qwen failed: {e}")
+                    final_prompt = original_prompt
+
+        # Generate seed
+        if seed_input == -1:
+            seed = random.randint(0, MAX_SEED)
+        else:
+            seed = int(seed_input)
         
+        generator = torch.Generator(device=device).manual_seed(seed)
+
         # Generate image
-        generator = torch.Generator(device="cuda").manual_seed(actual_seed)
-        
-        with torch.inference_mode():
-            image = pipeline(
+        progress(0.7, desc="Generating image...")
+        with torch.no_grad():
+            result = pipeline(
                 prompt=final_prompt,
+                height=height,
+                width=width,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
                 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!")
+        image = result.images[0]
         
-        status_log = f"Seed: {actual_seed} | Generation complete."
-        return image, final_prompt, status_log
+        # Clean up
+        del pipeline
+        torch.cuda.empty_cache()
         
+        progress(1.0, desc="Complete!")
+        return image, f"✅ Generation complete! Seed: {seed}", final_prompt
+
     except Exception as e:
         logging.error(f"Generation failed: {e}")
-        raise gr.Error(f"An error occurred: {str(e)}")
+        return None, f"❌ Generation failed: {str(e)}", ""
 
 # ------------------------------------------------------------------
 # 6. Gradio 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}")
+def create_interface():
+    """Create Gradio interface"""
     
-    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="Uses Qwen3-8B for prompt enhancement")
-            
-            with gr.Row():
-                width_input = gr.Slider(label="Width", minimum=256, maximum=2048, value=832, step=64)
-                height_input = gr.Slider(label="Height", minimum=256, maximum=2048, value=1216, step=64)
-            gr.Markdown("Tip: Recommended size is 832x1216 for best results.")
+    with gr.Blocks(
+        title="PosterCraft-v1.0",
+        theme=gr.themes.Soft(),
+        css="""
+        .main-container { max-width: 1200px; margin: 0 auto; }
+        .status-box { padding: 10px; border-radius: 5px; margin: 10px 0; }
+        .auth-success { background-color: #d4edda; border: 1px solid #c3e6cb; color: #155724; }
+        .auth-error { background-color: #f8d7da; border: 1px solid #f5c6cb; color: #721c24; }
+        """
+    ) as demo:
+        
+        gr.HTML("""
+        <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
+            </p>
+        </div>
+        """)
+        
+        with gr.Row():
+            gr.Markdown(f"**Base Pipeline:** `{DEFAULT_PIPELINE_PATH}`")
+            gr.Markdown(f"**Authentication Status:** {auth_status}")
+        
+        gr.HTML("""
+        <div class="status-box">
+            <p><strong>⚠️ First use requires model download, please wait about 10-15 minutes</strong></p>
+        </div>
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                original_prompt = gr.Textbox(
+                    label="Poster Prompt",
+                    placeholder="Enter your poster description...",
+                    lines=3,
+                    value="A vintage travel poster for Paris, featuring the Eiffel Tower at sunset with warm golden lighting"
+                )
+                
+                enable_recap = gr.Checkbox(
+                    label="Enable Prompt Enhancement (Qwen3-8B)",
+                    value=True,
+                    info="Use AI to enhance and expand your prompt"
+                )
+                
+                with gr.Row():
+                    height = gr.Slider(
+                        label="Height",
+                        minimum=256,
+                        maximum=MAX_IMAGE_SIZE,
+                        value=1024,
+                        step=32
+                    )
+                    width = gr.Slider(
+                        label="Width",
+                        minimum=256,
+                        maximum=MAX_IMAGE_SIZE,
+                        value=768,
+                        step=32
+                    )
+                
+                with gr.Row():
+                    num_inference_steps = gr.Slider(
+                        label="Inference Steps",
+                        minimum=1,
+                        maximum=50,
+                        value=20,
+                        step=1
+                    )
+                    guidance_scale = gr.Slider(
+                        label="Guidance Scale",
+                        minimum=1.0,
+                        maximum=15.0,
+                        value=3.5,
+                        step=0.1
+                    )
+                
+                seed_input = gr.Number(
+                    label="Seed (-1 for random)",
+                    value=-1,
+                    precision=0
+                )
+                
+                generate_btn = gr.Button(
+                    "🎨 Generate Poster",
+                    variant="primary",
+                    size="lg"
+                )
             
-            num_inference_steps_input = gr.Slider(label="Inference Steps", minimum=1, maximum=100, value=28, step=1)
-            guidance_scale_input = gr.Slider(label="Guidance Scale (CFG)", minimum=0.0, maximum=20.0, value=3.5, step=0.1)
-            seed_number_input = gr.Number(label="Seed", value=None, minimum=-1, step=1, info="Leave blank or set to -1 for a random seed.")
-            generate_button = gr.Button("Generate Image", variant="primary")
-
-        with gr.Column(scale=1):
-            gr.Markdown("### 2. Results")
-            image_output = gr.Image(label="Generated Image", type="pil", show_download_button=True, height=512)
-            recapped_prompt_output = gr.Textbox(label="Final Prompt Used", lines=5, interactive=False)
-            status_output = gr.Textbox(label="Status Log", lines=4, interactive=False)
-
-    inputs_list = [
-        prompt_input, enable_recap_checkbox, height_input, width_input,
-        num_inference_steps_input, guidance_scale_input, seed_number_input
-    ]
-    outputs_list = [image_output, recapped_prompt_output, status_output]
+            with gr.Column(scale=1):
+                output_image = gr.Image(
+                    label="Generated Poster",
+                    type="pil",
+                    height=600
+                )
+                
+                status_output = gr.Textbox(
+                    label="Generation Status",
+                    interactive=False,
+                    lines=2
+                )
+                
+                enhanced_prompt = gr.Textbox(
+                    label="Enhanced Prompt",
+                    interactive=False,
+                    lines=5,
+                    info="The final prompt used for generation"
+                )
+        
+        # Event handlers
+        generate_btn.click(
+            fn=generate_image_interface,
+            inputs=[
+                original_prompt, enable_recap, height, width,
+                num_inference_steps, guidance_scale, seed_input
+            ],
+            outputs=[output_image, status_output, enhanced_prompt]
+        )
+        
+        # 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=[original_prompt]
+        )
     
-    generate_button.click(fn=generate_image_interface, inputs=inputs_list, outputs=outputs_list)
+    return demo
 
+# ------------------------------------------------------------------
+# 7. Launch Application
+# ------------------------------------------------------------------
 if __name__ == "__main__":
-    demo.launch()
+    demo = create_interface()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_api=False
+    )