import gradio as gr
import numpy as np
import torch
import random
import os
import spaces
from PIL import Image, ImageOps, ImageFilter
from diffusers import FluxPipeline, DiffusionPipeline
import requests
from io import BytesIO
# Constants
MAX_SEED = np.iinfo(np.int32).max
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
print("Loading models...")
def load_pipeline():
"""Load the appropriate pipeline based on availability"""
global pipe, MODEL_STATUS
try:
# 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=512):
"""
Prepare reference and pose images for Kontext processing.
Following the RefControl format: reference (left) | pose (right)
"""
if reference_image is None or pose_image is None:
return None
# Convert to RGB
reference_image = reference_image.convert("RGB")
pose_image = pose_image.convert("RGB")
# Calculate dimensions maintaining aspect ratio
ref_ratio = reference_image.width / reference_image.height
pose_ratio = pose_image.width / pose_image.height
# Set heights to target size
height = target_size
ref_width = int(height * ref_ratio)
pose_width = int(height * pose_ratio)
# Ensure dimensions are divisible by 8 (FLUX requirement)
ref_width = (ref_width // 8) * 8
pose_width = (pose_width // 8) * 8
height = (height // 8) * 8
# Resize images
reference_resized = reference_image.resize((ref_width, height), Image.LANCZOS)
pose_resized = pose_image.resize((pose_width, height), Image.LANCZOS)
# Concatenate horizontally: reference | pose
total_width = ref_width + pose_width
concatenated = Image.new('RGB', (total_width, height))
concatenated.paste(reference_resized, (0, 0))
concatenated.paste(pose_resized, (ref_width, 0))
return concatenated
def process_pose_for_control(pose_image):
"""
Process pose image to ensure maximum contrast and clarity for control
"""
if pose_image is None:
return None
# Convert to grayscale first
gray = pose_image.convert("L")
# Apply strong edge detection
edges = gray.filter(ImageFilter.FIND_EDGES)
edges = edges.filter(ImageFilter.EDGE_ENHANCE_MORE)
# Maximize contrast
edges = ImageOps.autocontrast(edges, cutoff=2)
# Convert to pure black and white
threshold = 128
edges = edges.point(lambda x: 255 if x > threshold else 0, mode='1')
# Convert back to RGB with inverted colors (black lines on white)
edges = edges.convert("RGB")
edges = ImageOps.invert(edges)
return edges
@spaces.GPU(duration=60)
def generate_pose_transfer(
reference_image,
pose_image,
prompt="",
negative_prompt="",
seed=42,
randomize_seed=False,
guidance_scale=7.5, # Increased for better pose adherence
num_inference_steps=28,
lora_scale=1.0,
enhance_pose=False,
progress=gr.Progress(track_tqdm=True)
):
"""
Main generation function using RefControl approach.
"""
if pipe is None:
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")
# Randomize seed if requested
if randomize_seed:
seed = random.randint(0, MAX_SEED)
# Enhance pose if requested
if enhance_pose:
pose_image = process_pose_for_control(pose_image)
# Prepare concatenated input with fixed size
concatenated_input = prepare_images_for_kontext(reference_image, pose_image, target_size=512)
if concatenated_input is None:
raise gr.Error("Failed to process images")
# Ensure dimensions are model-compatible
width, height = concatenated_input.size
# Round to nearest 64 pixels for stability
width = (width // 64) * 64
height = (height // 64) * 64
# Limit maximum size to prevent memory issues
max_size = 1024
if width > max_size:
ratio = max_size / width
width = max_size
height = int(height * ratio)
height = (height // 64) * 64
if height > max_size:
ratio = max_size / height
height = max_size
width = int(width * ratio)
width = (width // 64) * 64
# Resize if needed
if (width, height) != concatenated_input.size:
concatenated_input = concatenated_input.resize((width, height), Image.LANCZOS)
# Construct prompt with trigger word - CRITICAL FOR POSE CONTROL
# The prompt must explicitly describe the pose transfer task
base_instruction = f"{TRIGGER_WORD}, A photo composed of two images side by side. Left: reference person. Right: target pose skeleton. Task: Generate the person from the left image in the exact pose shown in the right image"
if prompt:
full_prompt = f"{base_instruction}. Additional details: {prompt}"
else:
full_prompt = base_instruction
# Add strong pose control instructions
full_prompt += ". IMPORTANT: Strictly follow the pose/skeleton from the right image while preserving the identity, clothing, and appearance from the left image. The output should show ONLY the transformed person, not the side-by-side layout."
# Set generator for reproducibility
generator = torch.Generator("cuda").manual_seed(seed)
try:
# Check if we have LoRA capabilities
has_lora = hasattr(pipe, 'set_adapters') and "LoRA" in MODEL_STATUS
# Set LoRA with higher strength for better pose control
if has_lora:
try:
# Increase LoRA strength for pose control
actual_lora_scale = lora_scale * 1.5 # Boost LoRA influence
pipe.set_adapters(["refcontrol"], adapter_weights=[actual_lora_scale])
print(f"LoRA adapter set with boosted strength: {actual_lora_scale}")
except Exception as e:
print(f"LoRA adapter not set: {e}")
print(f"Generating with size: {width}x{height}")
print(f"Prompt: {full_prompt[:200]}...")
# Generate image with stronger pose control
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
if "Kontext" in MODEL_STATUS:
# Use Kontext pipeline - removed unsupported controlnet_conditioning_scale
result = pipe(
image=concatenated_input,
prompt=full_prompt,
negative_prompt=negative_prompt if negative_prompt else "blurry, distorted, deformed, wrong pose, incorrect posture",
guidance_scale=guidance_scale, # Higher for better control
num_inference_steps=num_inference_steps,
generator=generator,
width=width,
height=height,
).images[0]
else:
# Use standard FLUX pipeline
result = pipe(
prompt=full_prompt,
negative_prompt=negative_prompt if negative_prompt else "",
image=concatenated_input,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
generator=generator,
strength=0.85,
).images[0]
print("Generation successful!")
return result, seed, concatenated_input
except RuntimeError as e:
if "out of memory" in str(e).lower():
raise gr.Error("GPU out of memory. Try reducing image size or inference steps.")
else:
raise gr.Error(f"Generation failed: {str(e)}")
except Exception as e:
print(f"Error details: {e}")
raise gr.Error(f"Generation failed: {str(e)}")
# CSS styling
css = """
#col-container {
margin: 0 auto;
max-width: 1280px;
}
.header {
text-align: center;
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
padding: 20px;
border-radius: 12px;
margin-bottom: 20px;
}
.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;
border-radius: 8px;
overflow: hidden;
}
.result-image {
border: 3px solid #4CAF50;
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
gr.HTML("""
""")
# Model status
status_color = "#d4edda" if "✅" in MODEL_STATUS else "#fff3cd" if "⚠️" in MODEL_STATUS else "#f8d7da"
gr.HTML(f"""
{MODEL_STATUS}
""")
# Authentication check
if not HF_TOKEN:
gr.Markdown("""
### 🔐 Authentication Required
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("""
Note: For full LoRA support, PEFT library is required.
Add peft to your requirements.txt file.
""")
# Main interface
with gr.Row():
with gr.Column(scale=1):
gr.Markdown("### 📥 Input Images")
# Reference image
reference_image = gr.Image(
label="Reference Image (Subject to transform)",
type="pil",
elem_classes=["input-image"],
height=300
)
# Pose image
pose_image = gr.Image(
label="Pose Control (Line art or skeleton)",
type="pil",
elem_classes=["input-image"],
height=300
)
# Pose extraction tool
with gr.Accordion("🔧 Extract Pose from Image", open=False):
extract_source = gr.Image(
label="Source image for pose extraction",
type="pil",
height=200
)
extract_btn = gr.Button("Extract Pose", size="sm")
# Prompts
prompt = gr.Textbox(
label=f"Prompt ('{TRIGGER_WORD}' added automatically)",
placeholder="e.g., wearing elegant dress, professional photography",
lines=2
)
negative_prompt = gr.Textbox(
label="Negative Prompt (optional)",
placeholder="e.g., blurry, low quality, distorted",
lines=1,
value="blurry, low quality, distorted, deformed"
)
# Generate button
generate_btn = gr.Button(
"🎨 Generate Pose Transfer",
variant="primary",
size="lg"
)
# Advanced settings
with gr.Accordion("⚙️ Advanced Settings", open=False):
with gr.Row():
seed = gr.Slider(
label="Seed",
minimum=0,
maximum=MAX_SEED,
step=1,
value=42
)
randomize_seed = gr.Checkbox(
label="Randomize",
value=True
)
guidance_scale = gr.Slider(
label="Guidance Scale",
minimum=5.0,
maximum=15.0,
step=0.5,
value=7.5,
info="Higher = stricter pose following (7-10 recommended)"
)
num_inference_steps = gr.Slider(
label="Inference Steps",
minimum=20,
maximum=50,
step=1,
value=30
)
if "LoRA" in MODEL_STATUS:
lora_scale = gr.Slider(
label="LoRA Strength",
minimum=0.5,
maximum=2.0,
step=0.1,
value=1.2,
info="RefControl LoRA influence (1.0-1.5 recommended)"
)
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",
value=False
)
with gr.Column(scale=1):
gr.Markdown("### 🖼️ Result")
# Result image
result_image = gr.Image(
label="Generated Image",
elem_classes=["result-image"],
interactive=False,
height=500
)
# Seed display
seed_used = gr.Number(
label="Seed Used",
interactive=False
)
# Debug view
with gr.Accordion("🔍 Debug View", open=False):
concat_preview = gr.Image(
label="Input Concatenation (Reference | Pose)",
height=200
)
# Action buttons
with gr.Row():
reuse_ref_btn = gr.Button("♻️ Use as Reference", 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"],
["wearing red dress, outdoor garden"],
["business attire, office setting"],
["casual streetwear, urban background"],
["athletic wear, gym environment"],
],
inputs=[prompt]
)
# Instructions
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. **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
## Requirements:
- **HF_TOKEN**: Required for model access
- **peft**: Required for LoRA support (add to requirements.txt)
""")
# Event handlers
generate_btn.click(
fn=generate_pose_transfer,
inputs=[
reference_image,
pose_image,
prompt,
negative_prompt,
seed,
randomize_seed,
guidance_scale,
num_inference_steps,
lora_scale,
enhance_pose
],
outputs=[result_image, seed_used, concat_preview]
)
extract_btn.click(
fn=process_pose_for_control,
inputs=[extract_source],
outputs=[pose_image]
)
reuse_ref_btn.click(
fn=lambda x: x,
inputs=[result_image],
outputs=[reference_image]
)
reuse_pose_btn.click(
fn=process_pose_for_control,
inputs=[result_image],
outputs=[pose_image]
)
clear_btn.click(
fn=lambda: [None, None, "", "blurry, low quality, distorted, deformed", 42, None, None],
outputs=[
reference_image,
pose_image,
prompt,
negative_prompt,
seed_used,
result_image,
concat_preview
]
)
# Launch the app
if __name__ == "__main__":
demo.queue()
demo.launch()