Update app.py

app.py

@@ -6,7 +6,6 @@ import os
 import spaces
 from PIL import Image, ImageOps, ImageFilter
 from diffusers import FluxPipeline, DiffusionPipeline
-from diffusers.loaders import LoraLoaderMixin
 import requests
 from io import BytesIO
 
@@ -16,51 +15,70 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 
 # Model configuration
 KONTEXT_MODEL = "black-forest-labs/FLUX.1-Kontext-dev"
+FALLBACK_MODEL = "black-forest-labs/FLUX.1-dev"
 LORA_MODEL = "thedeoxen/refcontrol-flux-kontext-reference-pose-lora"
 TRIGGER_WORD = "refcontrolpose"
 
-# Initialize pipeline with authentication
+# Initialize pipeline
 print("Loading models...")
 
-try:
-    # Load Flux Kontext pipeline with HF token
-    if HF_TOKEN:
-        from diffusers import FluxKontextPipeline
-        pipe = FluxKontextPipeline.from_pretrained(
-            KONTEXT_MODEL,
-            torch_dtype=torch.bfloat16,
-            use_auth_token=HF_TOKEN
-        )
-        
-        # Load the RefControl LoRA
-        pipe.load_lora_weights(
-            LORA_MODEL,
-            adapter_name="refcontrol",
-            use_auth_token=HF_TOKEN
-        )
-        
-        # Move to GPU
-        pipe = pipe.to("cuda")
-        
-        MODEL_STATUS = "✅ Flux Kontext + RefControl LoRA loaded successfully"
-        print(MODEL_STATUS)
-        
-    else:
-        raise ValueError("HF_TOKEN not found in environment variables")
-        
-except Exception as e:
-    print(f"Error loading models: {e}")
-    # Fallback to base model without LoRA
+def load_pipeline():
+    """Load the appropriate pipeline based on availability"""
+    global pipe, MODEL_STATUS
+    
     try:
-        pipe = DiffusionPipeline.from_pretrained(
-            "black-forest-labs/FLUX.1-dev",
-            torch_dtype=torch.bfloat16,
-            use_auth_token=HF_TOKEN if HF_TOKEN else True
-        ).to("cuda")
-        MODEL_STATUS = "⚠️ Running in fallback mode (FLUX.1-dev without LoRA)"
-    except:
-        MODEL_STATUS = "❌ Failed to load models. Please check HF_TOKEN"
+        # First, try to import necessary libraries
+        try:
+            from diffusers import FluxKontextPipeline
+            import peft
+            print("PEFT library found")
+            use_kontext = True
+        except ImportError:
+            print("FluxKontextPipeline or PEFT not available, using fallback")
+            use_kontext = False
+        
+        if use_kontext and HF_TOKEN:
+            # Try to load Kontext model
+            pipe = FluxKontextPipeline.from_pretrained(
+                KONTEXT_MODEL,
+                torch_dtype=torch.bfloat16,
+                token=HF_TOKEN
+            )
+            
+            # Try to load LoRA if PEFT is available
+            try:
+                pipe.load_lora_weights(
+                    LORA_MODEL,
+                    adapter_name="refcontrol",
+                    token=HF_TOKEN
+                )
+                MODEL_STATUS = "✅ Flux Kontext + RefControl LoRA loaded"
+            except Exception as e:
+                print(f"Could not load LoRA: {e}")
+                MODEL_STATUS = "⚠️ Flux Kontext loaded (without LoRA - PEFT required)"
+            
+            pipe = pipe.to("cuda")
+            
+        else:
+            # Fallback to standard FLUX
+            pipe = FluxPipeline.from_pretrained(
+                FALLBACK_MODEL,
+                torch_dtype=torch.bfloat16,
+                token=HF_TOKEN if HF_TOKEN else True
+            )
+            pipe = pipe.to("cuda")
+            MODEL_STATUS = "⚠️ Using FLUX.1-dev (fallback mode)"
+            
+    except Exception as e:
+        print(f"Error loading models: {e}")
+        MODEL_STATUS = f"❌ Error: {str(e)}"
         pipe = None
+    
+    return pipe, MODEL_STATUS
+
+# Load the pipeline
+pipe, MODEL_STATUS = load_pipeline()
+print(MODEL_STATUS)
 
 def prepare_images_for_kontext(reference_image, pose_image, target_size=768):
     """
@@ -140,11 +158,11 @@ def generate_pose_transfer(
     progress=gr.Progress(track_tqdm=True)
 ):
     """
-    Main generation function using RefControl LoRA.
+    Main generation function using RefControl approach.
     """
     
     if pipe is None:
-        return None, 0, "Model not loaded. Please check HF_TOKEN"
+        return None, 0, "Model not loaded. Please check HF_TOKEN and restart the Space"
     
     if reference_image is None or pose_image is None:
         raise gr.Error("Please upload both reference and pose images")
@@ -176,23 +194,40 @@ def generate_pose_transfer(
     generator = torch.Generator("cuda").manual_seed(seed)
     
     try:
-        # Generate with LoRA
+        # Check if we have LoRA capabilities
+        has_lora = hasattr(pipe, 'set_adapters') and "RefControl" in MODEL_STATUS
+        
         with torch.autocast("cuda"):
-            if hasattr(pipe, 'set_adapters'):
-                # Set LoRA adapter strength
-                pipe.set_adapters(["refcontrol"], adapter_weights=[lora_scale])
+            if has_lora:
+                # Try to set LoRA adapter strength
+                try:
+                    pipe.set_adapters(["refcontrol"], adapter_weights=[lora_scale])
+                except Exception as e:
+                    print(f"Could not set LoRA adapter: {e}")
             
-            # Generate image
-            result = pipe(
-                image=concatenated_input,
-                prompt=full_prompt,
-                negative_prompt=negative_prompt,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                generator=generator,
-                width=concatenated_input.width,
-                height=concatenated_input.height,
-            ).images[0]
+            # Generate image based on pipeline type
+            if "Kontext" in MODEL_STATUS:
+                # Use Kontext pipeline
+                result = pipe(
+                    image=concatenated_input,
+                    prompt=full_prompt,
+                    negative_prompt=negative_prompt if negative_prompt else None,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_inference_steps,
+                    generator=generator,
+                    width=concatenated_input.width,
+                    height=concatenated_input.height,
+                ).images[0]
+            else:
+                # Use standard FLUX pipeline (image-to-image)
+                result = pipe(
+                    prompt=full_prompt,
+                    image=concatenated_input,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_inference_steps,
+                    generator=generator,
+                    strength=0.85,  # For img2img mode
+                ).images[0]
             
         return result, seed, concatenated_input
         
@@ -215,12 +250,14 @@ css = """
 .header h1 {
     color: white;
     margin: 0;
+    font-size: 2em;
 }
 .status-box {
     padding: 10px;
     border-radius: 8px;
     margin: 10px 0;
     font-weight: bold;
+    text-align: center;
 }
 .input-image {
     border: 2px solid #e0e0e0;
@@ -232,40 +269,57 @@ css = """
     border-radius: 8px;
     overflow: hidden;
 }
+.info-box {
+    background: #f0f0f0;
+    padding: 10px;
+    border-radius: 8px;
+    margin: 10px 0;
+}
 """
 
 # Create Gradio interface
 with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     
     with gr.Column(elem_id="col-container"):
-        # Header with authentication
-        with gr.Row():
-            with gr.Column():
-                gr.HTML("""
-                <div class="header">
-                    <h1>🎭 RefControl Flux Kontext - Reference Pose Transfer</h1>
-                    <p style="color: white;">Powered by thedeoxen/refcontrol-flux-kontext-reference-pose-lora</p>
-                </div>
-                """)
-                
-                # Model status
-                gr.Markdown(f"""
-                <div class="status-box" style="background: {'#d4edda' if '✅' in MODEL_STATUS else '#f8d7da'};">
-                    {MODEL_STATUS}
-                </div>
-                """)
+        # Header
+        gr.HTML("""
+        <div class="header">
+            <h1>🎭 FLUX Pose Transfer System</h1>
+            <p style="color: white;">Transfer poses while preserving identity</p>
+        </div>
+        """)
+        
+        # Model status
+        status_color = "#d4edda" if "✅" in MODEL_STATUS else "#fff3cd" if "⚠️" in MODEL_STATUS else "#f8d7da"
+        gr.HTML(f"""
+        <div class="status-box" style="background: {status_color};">
+            {MODEL_STATUS}
+        </div>
+        """)
         
-        # Authentication info
+        # Authentication check
         if not HF_TOKEN:
             gr.Markdown("""
             ### 🔐 Authentication Required
-            Please set your Hugging Face token to use this Space:
-            1. Go to Settings → Variables and secrets
+            
+            To use this Space with full features:
+            1. Go to **Settings** → **Variables and secrets**
             2. Add `HF_TOKEN` with your Hugging Face token
             3. Restart the Space
+            
+            Or click below to sign in:
             """)
             gr.LoginButton("Sign in with Hugging Face", size="lg")
         
+        # Info box for PEFT requirement
+        if "PEFT required" in MODEL_STATUS:
+            gr.HTML("""
+            <div class="info-box">
+                <b>Note:</b> For full LoRA support, PEFT library is required. 
+                Add <code>peft</code> to your requirements.txt file.
+            </div>
+            """)
+        
         # Main interface
         with gr.Row():
             with gr.Column(scale=1):
@@ -298,15 +352,16 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                 
                 # Prompts
                 prompt = gr.Textbox(
-                    label=f"Prompt (trigger word '{TRIGGER_WORD}' added automatically)",
+                    label=f"Prompt ('{TRIGGER_WORD}' added automatically)",
                     placeholder="e.g., wearing elegant dress, professional photography",
                     lines=2
                 )
                 
                 negative_prompt = gr.Textbox(
-                    label="Negative Prompt",
+                    label="Negative Prompt (optional)",
                     placeholder="e.g., blurry, low quality, distorted",
-                    lines=2
+                    lines=1,
+                    value="blurry, low quality, distorted, deformed"
                 )
                 
                 # Generate button
@@ -348,14 +403,24 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                         value=28
                     )
                     
-                    lora_scale = gr.Slider(
-                        label="LoRA Strength",
-                        minimum=0.0,
-                        maximum=2.0,
-                        step=0.1,
-                        value=1.0,
-                        info="RefControl LoRA influence"
-                    )
+                    if "LoRA" in MODEL_STATUS:
+                        lora_scale = gr.Slider(
+                            label="LoRA Strength",
+                            minimum=0.0,
+                            maximum=2.0,
+                            step=0.1,
+                            value=1.0,
+                            info="RefControl LoRA influence"
+                        )
+                    else:
+                        lora_scale = gr.Slider(
+                            label="LoRA Strength (not available)",
+                            minimum=0.0,
+                            maximum=2.0,
+                            step=0.1,
+                            value=1.0,
+                            interactive=False
+                        )
                     
                     enhance_pose = gr.Checkbox(
                         label="Auto-enhance pose edges",
@@ -363,7 +428,7 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                     )
             
             with gr.Column(scale=1):
-                gr.Markdown("### 🖼️ Generated Result")
+                gr.Markdown("### 🖼️ Result")
                 
                 # Result image
                 result_image = gr.Image(
@@ -373,12 +438,11 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                     height=500
                 )
                 
-                # Info display
-                with gr.Row():
-                    seed_used = gr.Number(
-                        label="Seed Used",
-                        interactive=False
-                    )
+                # Seed display
+                seed_used = gr.Number(
+                    label="Seed Used",
+                    interactive=False
+                )
                 
                 # Debug view
                 with gr.Accordion("🔍 Debug View", open=False):
@@ -387,39 +451,48 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                         height=200
                     )
                 
-                # Reuse buttons
+                # Action buttons
                 with gr.Row():
                     reuse_ref_btn = gr.Button("♻️ Use as Reference", size="sm")
-                    reuse_pose_btn = gr.Button("📐 Extract & Use as Pose", size="sm")
+                    reuse_pose_btn = gr.Button("📐 Extract Pose", size="sm")
                     clear_btn = gr.Button("🗑️ Clear All", size="sm")
         
         # Examples
         gr.Markdown("### 💡 Example Prompts")
         gr.Examples(
             examples=[
-                ["professional portrait, studio lighting, high quality"],
-                ["wearing red dress, outdoor garden setting"],
-                ["business attire, corporate headshot"],
+                ["professional portrait, studio lighting"],
+                ["wearing red dress, outdoor garden"],
+                ["business attire, office setting"],
                 ["casual streetwear, urban background"],
-                ["athletic wear, dynamic action shot"],
-                ["elegant evening gown, luxury setting"],
+                ["athletic wear, gym environment"],
             ],
             inputs=[prompt]
         )
         
         # Instructions
-        with gr.Accordion("📖 How to Use", open=False):
-            gr.Markdown("""
-            1. **Upload Reference Image**: The person/subject you want to transform
+        with gr.Accordion("📖 Instructions", open=False):
+            gr.Markdown(f"""
+            ## How to Use:
+            
+            1. **Upload Reference Image**: The person whose appearance you want to keep
             2. **Upload Pose Image**: Line art or skeleton pose to follow
-            3. **Optional**: Add descriptive prompt for style/setting
-            4. **Click Generate**: Wait for the magic to happen!
+            3. **Add Prompt** (optional): Describe additional details
+            4. **Click Generate**: Create your pose-transferred image
+            
+            ## Model Information:
+            - **Current Model**: {MODEL_STATUS}
+            - **Trigger Word**: `{TRIGGER_WORD}` (added automatically)
+            
+            ## Tips:
+            - Use clear, high-contrast pose images
+            - Black lines on white background work best for poses
+            - Adjust guidance scale for pose adherence strength
+            - Higher steps = better quality but slower
             
-            **Tips:**
-            - Use clear, high-contrast pose images for best results
-            - The model preserves identity from reference while following pose
-            - Adjust LoRA strength to balance identity vs pose adherence
-            - Higher guidance scale = stricter pose following
+            ## Requirements:
+            - **HF_TOKEN**: Required for model access
+            - **peft**: Required for LoRA support (add to requirements.txt)
             """)
     
     # Event handlers
@@ -459,7 +532,7 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     )
     
     clear_btn.click(
-        fn=lambda: [None, None, "", "", 42, None, None],
+        fn=lambda: [None, None, "", "blurry, low quality, distorted, deformed", 42, None, None],
         outputs=[
             reference_image,
             pose_image,