app.py

import os
import torch
import gc
import gradio as gr
import numpy as np
from PIL import Image
from einops import rearrange
import io
import requests
import spaces
from huggingface_hub import login
from gradio_imageslider import ImageSlider
from diffusers.utils import load_image
from diffusers import FluxControlNetPipeline, FluxControlNetModel

# Device settings: CPU for loading, GPU for inference
device_cpu = torch.device("cpu")
device_gpu = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Model identifiers
base_model = 'black-forest-labs/FLUX.1-dev'
controlnet_model = 'InstantX/FLUX.1-dev-Controlnet-Union'

# Load the ControlNet model and pipeline on CPU
controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16).to(device_cpu)
pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16).to(device_cpu)

controlnet_conditioning_scale = 0.5

control_modes = {
    "canny": 0,
    "tile": 1,
    "depth": 2,
    "blur": 3,
    "pose": 4,
    "gray": 5,
    "lq": 6,
}

def load_and_convert_image(image):
    """Load and convert images to a format that PIL can handle."""
    if isinstance(image, str):
        image = Image.open(image)
    elif isinstance(image, bytes):
        image = Image.open(io.BytesIO(image))
    # Convert AVIF to PNG if necessary
    if image.format == 'AVIF':
        image = image.convert("RGB")  # Convert to a format PIL can handle
    return image

def preprocess_image(image, target_width, target_height, crop=True):
    """Preprocess image to match the target dimensions."""
    image = load_and_convert_image(image)
    if crop:
        original_width, original_height = image.size

        # Resize to match the target size without stretching
        scale = max(target_width / original_width, target_height / original_height)
        resized_width = int(scale * original_width)
        resized_height = int(scale * original_height)

        image = image.resize((resized_width, resized_height), Image.LANCZOS)
        
        # Center crop to match the target dimensions
        left = (resized_width - target_width) // 2
        top = (resized_height - target_height) // 2
        image = image.crop((left, top, left + target_width, top + target_height))
    else:
        image = image.resize((target_width, target_height), Image.LANCZOS)
    
    return image

def clear_cuda_memory():
    """Clear CUDA memory."""
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()

@spaces.GPU(duration=120)
def generate_image(prompt, control_image, control_mode, num_steps=50, guidance=4, width=512, height=512, seed=42, random_seed=False):
    """Generate image using the FLUX.1 ControlNet model."""
    clear_cuda_memory()
    
    if random_seed:
        seed = np.random.randint(0, 10000)
    
    if not os.path.isdir("./controlnet_results/"):
        os.makedirs("./controlnet_results/")

    # Move model to GPU for inference
    pipe.to(device_gpu)

    control_image = preprocess_image(control_image, width, height)
    
    torch.manual_seed(seed)
    with torch.no_grad():
        image = pipe(
            prompt,
            control_image=control_image,
            control_mode=control_modes[control_mode],
            width=width,
            height=height,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            num_inference_steps=num_steps,
            guidance_scale=guidance,
        ).images[0]
    
    # Move model back to CPU after inference
    pipe.to(device_cpu)

    return [control_image, image]  # Return both images for slider

interface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Image(type="pil", label="Control Image"),
        gr.Dropdown(choices=list(control_modes.keys()), label="Control Mode", value="canny"),
        gr.Slider(step=1, minimum=1, maximum=64, value=28, label="Num Steps"),
        gr.Slider(minimum=0.1, maximum=10, value=4, label="Guidance"),
        gr.Slider(minimum=128, maximum=2048, step=128, value=1024, label="Width"),
        gr.Slider(minimum=128, maximum=2048, step=128, value=1024, label="Height"),
        gr.Number(value=42, label="Seed"),
        gr.Checkbox(label="Random Seed")
    ],
    outputs=ImageSlider(label="Before / After"),  # Use ImageSlider as the output
    title="FLUX.1 Controlnet Canny",
    description="Generate images using ControlNet and a text prompt.\n[[non-commercial license, Flux.1 Dev](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)]"
)

if __name__ == "__main__":
    interface.launch(share=True)