app.py

import sys
sys.path.append('./')

import gradio as gr
import spaces
import os
import sys
import subprocess
import numpy as np
from PIL import Image
import cv2
import torch
import random
from transformers import pipeline

# Skip trying to install the extension since it's failing
# We'll implement the necessary functions directly
print("Skipping ControlNet annotator installation - will use built-in implementations")

# Simplified translation function that just passes through text
# since the translation models are causing issues
def translate_to_english(text):
    # Check if Korean characters are present
    if any('\uAC00' <= char <= '\uD7A3' for char in text):
        print(f"Korean text detected: {text}")
        print("Translation is disabled - using original text")
    return text

from huggingface_hub import hf_hub_download

from huggingface_hub import login
hf_token = os.environ.get("HF_TOKEN_GATED")
login(token=hf_token)

MAX_SEED = np.iinfo(np.int32).max

def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed

# Define our own implementations since the imports are failing

# Simple Canny edge detector class
class CannyDetector:
    def __call__(self, image, low_threshold=100, high_threshold=200):
        # Convert PIL Image to cv2
        img = np.array(image)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        
        # Apply Canny edge detection
        canny = cv2.Canny(img, low_threshold, high_threshold)
        canny = cv2.dilate(canny, np.ones((2, 2), np.uint8), iterations=1)
        
        # Convert back to PIL
        return Image.fromarray(canny).convert("RGB")

# Simple OpenPose detector (placeholder implementation)
class OpenposeDetector:
    @classmethod
    def from_pretrained(cls, model_path):
        return cls()
        
    def __call__(self, image, hand_and_face=True):
        # For now, just use a basic person detection
        # In a real implementation, this would perform actual pose estimation
        # Here we're just creating a simple representation of a person
        
        # Create a white canvas of the same size as input
        img = np.array(image)
        h, w = img.shape[:2]
        canvas = np.ones((h, w, 3), dtype=np.uint8) * 255
        
        # Draw a simple stick figure in the center
        center_x, center_y = w//2, h//2
        head_radius = min(h, w) // 10
        body_length = head_radius * 4
        
        # Head
        cv2.circle(canvas, (center_x, center_y - head_radius), head_radius, (0, 0, 255), 2)
        
        # Body
        cv2.line(canvas, (center_x, center_y), (center_x, center_y + body_length), (0, 0, 255), 2)
        
        # Arms
        cv2.line(canvas, (center_x, center_y + head_radius), 
                (center_x - head_radius*2, center_y + head_radius*2), (0, 0, 255), 2)
        cv2.line(canvas, (center_x, center_y + head_radius), 
                (center_x + head_radius*2, center_y + head_radius*2), (0, 0, 255), 2)
        
        # Legs
        cv2.line(canvas, (center_x, center_y + body_length), 
                (center_x - head_radius*1.5, center_y + body_length + head_radius*3), (0, 0, 255), 2)
        cv2.line(canvas, (center_x, center_y + body_length), 
                (center_x + head_radius*1.5, center_y + body_length + head_radius*3), (0, 0, 255), 2)
        
        return Image.fromarray(canvas)

from depth_anything_v2.dpt import DepthAnythingV2

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
model_configs = {
    'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
    'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
    'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
    'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]}
}

encoder = 'vitl'
model = DepthAnythingV2(**model_configs[encoder])
filepath = hf_hub_download(repo_id=f"depth-anything/Depth-Anything-V2-Large", filename=f"depth_anything_v2_vitl.pth", repo_type="model")
state_dict = torch.load(filepath, map_location="cpu")
model.load_state_dict(state_dict)
model = model.to(DEVICE).eval()

import torch
from diffusers.utils import load_image
from diffusers import FluxControlNetPipeline, FluxControlNetModel
from diffusers.models import FluxMultiControlNetModel

base_model = 'black-forest-labs/FLUX.1-dev'
controlnet_model = 'Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro'
controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
controlnet = FluxMultiControlNetModel([controlnet])
pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16)
pipe.to("cuda")

# Fixed dictionary keys to use English for consistency
mode_mapping = {"Canny":0, "Tile":1, "Depth":2, "Blur":3, "OpenPose":4, "Grayscale":5, "LowQuality": 6}
strength_mapping = {"Canny":0.65, "Tile":0.45, "Depth":0.55, "Blur":0.45, "OpenPose":0.55, "Grayscale":0.45, "LowQuality": 0.4}

# Use our custom detector classes
canny = CannyDetector()
open_pose = OpenposeDetector.from_pretrained("lllyasviel/Annotators")

torch.backends.cuda.matmul.allow_tf32 = True
pipe.vae.enable_tiling()
pipe.vae.enable_slicing()
pipe.enable_model_cpu_offload() # for saving memory

def convert_from_image_to_cv2(img: Image) -> np.ndarray:
    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

def convert_from_cv2_to_image(img: np.ndarray) -> Image:
    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))

def extract_depth(image):
    image = np.asarray(image)
    depth = model.infer_image(image[:, :, ::-1])
    depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
    depth = depth.astype(np.uint8)
    gray_depth = Image.fromarray(depth).convert('RGB') 
    return gray_depth

def extract_openpose(img):
    processed_image_open_pose = open_pose(img, hand_and_face=True)
    return processed_image_open_pose
    
def extract_canny(image):
    processed_image_canny = canny(image)
    return processed_image_canny

def apply_gaussian_blur(image, kernel_size=(21, 21)):
    image = convert_from_image_to_cv2(image)
    blurred_image = convert_from_cv2_to_image(cv2.GaussianBlur(image, kernel_size, 0))
    return blurred_image

def convert_to_grayscale(image):
    image = convert_from_image_to_cv2(image)
    gray_image = convert_from_cv2_to_image(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY))
    return gray_image

def add_gaussian_noise(image, mean=0, sigma=10):
    image = convert_from_image_to_cv2(image)
    noise = np.random.normal(mean, sigma, image.shape)
    noisy_image = convert_from_cv2_to_image(np.clip(image.astype(np.float32) + noise, 0, 255).astype(np.uint8))
    return noisy_image

def tile(input_image, resolution=768):
    input_image = convert_from_image_to_cv2(input_image)
    H, W, C = input_image.shape
    H = float(H)
    W = float(W)
    k = float(resolution) / min(H, W)
    H *= k
    W *= k
    H = int(np.round(H / 64.0)) * 64
    W = int(np.round(W / 64.0)) * 64
    img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
    img = convert_from_cv2_to_image(img)
    return img

def resize_img(input_image, max_side=768, min_side=512, size=None, 
               pad_to_max_side=False, mode=Image.BILINEAR, base_pixel_number=64):

    w, h = input_image.size
    if size is not None:
        w_resize_new, h_resize_new = size
    else:
        ratio = min_side / min(h, w)
        w, h = round(ratio*w), round(ratio*h)
        ratio = max_side / max(h, w)
        input_image = input_image.resize([round(ratio*w), round(ratio*h)], mode)
        w_resize_new = (round(ratio * w) // base_pixel_number) * base_pixel_number
        h_resize_new = (round(ratio * h) // base_pixel_number) * base_pixel_number
    input_image = input_image.resize([w_resize_new, h_resize_new], mode)

    if pad_to_max_side:
        res = np.ones([max_side, max_side, 3], dtype=np.uint8) * 255
        offset_x = (max_side - w_resize_new) // 2
        offset_y = (max_side - h_resize_new) // 2
        res[offset_y:offset_y+h_resize_new, offset_x:offset_x+w_resize_new] = np.array(input_image)
        input_image = Image.fromarray(res)
    return input_image

@spaces.GPU()
def infer(cond_in, image_in, prompt, inference_steps, guidance_scale, control_mode, control_strength, seed, progress=gr.Progress(track_tqdm=True)):
    try:    
        control_mode_num = mode_mapping[control_mode]
        prompt = translate_to_english(prompt)
        
        if cond_in is None:
            if image_in is not None:
                image_in = resize_img(load_image(image_in))
                if control_mode == "Canny":
                    control_image = extract_canny(image_in)
                elif control_mode == "Depth":
                    control_image = extract_depth(image_in)
                elif control_mode == "OpenPose":
                    control_image = extract_openpose(image_in)
                elif control_mode == "Blur":
                    control_image = apply_gaussian_blur(image_in)
                elif control_mode == "LowQuality":
                    control_image = add_gaussian_noise(image_in)
                elif control_mode == "Grayscale":
                    control_image = convert_to_grayscale(image_in)
                elif control_mode == "Tile":
                    control_image = tile(image_in)
        else:
            control_image = resize_img(load_image(cond_in))

        width, height = control_image.size
        
        image = pipe(
            prompt, 
            control_image=[control_image],
            control_mode=[control_mode_num],
            width=width,
            height=height,
            controlnet_conditioning_scale=[control_strength],
            num_inference_steps=inference_steps, 
            guidance_scale=guidance_scale,
            generator=torch.manual_seed(seed),
        ).images[0]

        torch.cuda.empty_cache() 
        
        return image, control_image, gr.update(visible=True)
    
    except Exception as e:
        print(f"Error in inference: {e}")
        return None, None, gr.update(visible=True)
   

css = """
footer {
    visibility: hidden;
}
"""

with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
    with gr.Column(elem_id="col-container"):
        
        with gr.Column():
            
            with gr.Row():
                with gr.Column():
                    
                    with gr.Row(equal_height=True):
                        cond_in = gr.Image(label="Upload Processed Control Image", sources=["upload"], type="filepath")
                        image_in = gr.Image(label="Extract Condition from Reference Image (Optional)", sources=["upload"], type="filepath")
                    
                    prompt = gr.Textbox(label="Prompt", value="Highest Quality")
                    
                    with gr.Accordion("ControlNet"):
                        control_mode = gr.Radio(
                            ["Canny", "Depth", "OpenPose", "Grayscale", "Blur", "Tile", "LowQuality"], 
                            label="Mode", 
                            value="Grayscale",
                            info="Select control mode, applies to all images"
                        )
                        
                        control_strength = gr.Slider(
                            label="Control Strength",
                            minimum=0,
                            maximum=1.0,
                            step=0.05,
                            value=0.50,
                        )
                    
                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=42,
                    )
                    randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                    
                    with gr.Accordion("Advanced Settings", open=False):
                        with gr.Column():
                            with gr.Row():
                                inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=24)
                                guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=3.5)
                    
                    submit_btn = gr.Button("Submit")
                    
                with gr.Column():
                    result = gr.Image(label="Result")
                    processed_cond = gr.Image(label="Preprocessed Condition")

    submit_btn.click(
        fn=randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False
    ).then(
        fn = infer,
        inputs = [cond_in, image_in, prompt, inference_steps, guidance_scale, control_mode, control_strength, seed],
        outputs = [result, processed_cond],
        show_api=False
    )

demo.queue(api_open=False)
demo.launch()