app.py

import logging
import random
import warnings
import os
import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import FluxControlNetModel
from diffusers.pipelines import FluxControlNetPipeline
from gradio_imageslider import ImageSlider
from PIL import Image
from huggingface_hub import snapshot_download

css = """
#col-container {
    margin: 0 auto;
    max-width: 512px;
}
"""

# Device and dtype setup with lower precision
if torch.cuda.is_available():
    power_device = "GPU"
    device = "cuda" 
    dtype = torch.float16  # Changed to float16 for less memory usage
else:
    power_device = "CPU"
    device = "cpu"
    dtype = torch.float32

# Reduce CUDA memory usage
torch.cuda.empty_cache()
if torch.cuda.is_available():
    torch.cuda.set_per_process_memory_fraction(0.7)  # Use only 70% of GPU memory

huggingface_token = os.getenv("HUGGINFACE_TOKEN")

model_path = snapshot_download(
    repo_id="black-forest-labs/FLUX.1-dev",
    repo_type="model",
    ignore_patterns=["*.md", "*..gitattributes"],
    local_dir="FLUX.1-dev", 
    token=huggingface_token,
)

# Load pipeline with more memory optimizations
controlnet = FluxControlNetModel.from_pretrained(
    "jasperai/Flux.1-dev-Controlnet-Upscaler",
    torch_dtype=dtype,
    low_cpu_mem_usage=True,
    use_safetensors=True
).to(device)

pipe = FluxControlNetPipeline.from_pretrained(
    model_path,
    controlnet=controlnet,
    torch_dtype=dtype,
    low_cpu_mem_usage=True,
    use_safetensors=True
)

# Enable all possible memory optimizations
pipe.enable_model_cpu_offload()
pipe.enable_attention_slicing(1)
pipe.enable_sequential_cpu_offload()
pipe.enable_vae_slicing()

# Further reduce memory usage
MAX_SEED = 1000000
MAX_PIXEL_BUDGET = 256 * 256  # Further reduced from 512 * 512

def check_resources():
    if torch.cuda.is_available():
        gpu_memory = torch.cuda.get_device_properties(0).total_memory
        memory_allocated = torch.cuda.memory_allocated(0)
        if memory_allocated/gpu_memory > 0.8:  # 80% threshold
            return False
    return True

def process_input(input_image, upscale_factor, **kwargs):
    # Convert image to RGB mode to ensure compatibility
    input_image = input_image.convert('RGB')
    
    w, h = input_image.size
    w_original, h_original = w, h
    aspect_ratio = w / h

    was_resized = False

    if w * h * upscale_factor**2 > MAX_PIXEL_BUDGET:
        warnings.warn(
            f"Requested output image is too large. Resizing..."
        )
        gr.Info(
            f"Resizing input image to fit memory constraints..."
        )
        input_image = input_image.resize(
            (
                int(aspect_ratio * MAX_PIXEL_BUDGET**0.5 // upscale_factor),
                int(MAX_PIXEL_BUDGET**0.5 // aspect_ratio // upscale_factor),
            ),
            Image.LANCZOS
        )
        was_resized = True

    # resize to multiple of 8
    w, h = input_image.size
    w = w - w % 8
    h = h - h % 8

    return input_image.resize((w, h)), w_original, h_original, was_resized

@spaces.GPU
def infer(
    seed,
    randomize_seed,
    input_image,
    num_inference_steps,
    upscale_factor,
    controlnet_conditioning_scale,
    progress=gr.Progress(track_tqdm=True),
):
    try:
        if not check_resources():
            gr.Warning("System resources are running low. Try reducing parameters.")
            return None

        # Clear CUDA cache before processing
        if device == "cuda":
            torch.cuda.empty_cache()

        if randomize_seed:
            seed = random.randint(0, MAX_SEED)
            
        true_input_image = input_image
        input_image, w_original, h_original, was_resized = process_input(
            input_image, upscale_factor
        )

        # rescale with upscale factor
        w, h = input_image.size
        control_image = input_image.resize((w * upscale_factor, h * upscale_factor))

        generator = torch.Generator().manual_seed(seed)

        gr.Info("Upscaling image...")
        with torch.inference_mode():  # Use inference mode to save memory
            image = pipe(
                prompt="",
                control_image=control_image,
                controlnet_conditioning_scale=controlnet_conditioning_scale,
                num_inference_steps=num_inference_steps,
                guidance_scale=3.5,
                height=control_image.size[1],
                width=control_image.size[0],
                generator=generator,
            ).images[0]

        if was_resized:
            gr.Info(
                f"Resizing output image to final size..."
            )

        # resize to target desired size
        image = image.resize((w_original * upscale_factor, h_original * upscale_factor))
        return [true_input_image, image, seed]

    except RuntimeError as e:
        if "out of memory" in str(e):
            gr.Warning("Not enough GPU memory. Try reducing the upscale factor or image size.")
            return None
        raise e
    except Exception as e:
        gr.Error(f"An error occurred: {str(e)}")
        return None

with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
    with gr.Row():
        run_button = gr.Button(value="Run")

    with gr.Row():
        with gr.Column(scale=4):
            input_im = gr.Image(label="Input Image", type="pil")
        with gr.Column(scale=1):
            num_inference_steps = gr.Slider(
                label="Number of Inference Steps",
                minimum=8,
                maximum=30,  # Reduced from 50
                step=1,
                value=20,    # Reduced from 28
            )
            upscale_factor = gr.Slider(
                label="Upscale Factor",
                minimum=1,
                maximum=2,
                step=1,
                value=1,    # Reduced default
            )
            controlnet_conditioning_scale = gr.Slider(
                label="Controlnet Conditioning Scale",
                minimum=0.1,
                maximum=1.0,  # Reduced from 1.5
                step=0.1,
                value=0.5,   # Reduced from 0.6
            )
            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.Row():
        result = ImageSlider(label="Input / Output", type="pil", interactive=True)

    current_dir = os.path.dirname(os.path.abspath(__file__))
    
    examples = gr.Examples(
        examples=[
            [42, False, os.path.join(current_dir, "z1.webp"), 20, 1, 0.5],  # Reduced parameters
            [42, False, os.path.join(current_dir, "z2.webp"), 20, 1, 0.5],  # Reduced parameters
        ],
        inputs=[
            seed,
            randomize_seed,
            input_im,
            num_inference_steps,
            upscale_factor,
            controlnet_conditioning_scale,
        ],
        fn=infer,
        outputs=result,
        cache_examples="lazy",
    )

    gr.on(
        [run_button.click],
        fn=infer,
        inputs=[
            seed,
            randomize_seed,
            input_im,
            num_inference_steps,
            upscale_factor,
            controlnet_conditioning_scale,
        ],
        outputs=result,
        show_api=False,
    )

# Launch with minimal memory usage
demo.queue(max_size=1).launch(
    share=False,
    debug=True,
    show_error=True,
    max_threads=1,
    enable_queue=True
)