app.py

import torch
import numpy as np
import gradio as gr
from PIL import Image
import math
import torch.nn.functional as F
import os
import tempfile
import time
import threading

from utils.hatropeamp import HATNOUP_ROPE_AMP
from utils.fea2gsropeamp import Fea2GS_ROPE_AMP
from utils.edsrbaseline import EDSRNOUP
from utils.hatropeamp import HATNOUP_ROPE_AMP
from utils.rdn import RDNNOUP
from utils.swinir import SwinIRNOUP
from utils.fea2gsropeamp import Fea2GS_ROPE_AMP
from utils.gaussian_splatting import generate_2D_gaussian_splatting_step
from utils.split_and_joint_image import split_and_joint_image
from huggingface_hub import hf_hub_download
import subprocess
import sys
import spaces



# Device setup
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Global stop flag for interrupting inference
stop_inference = False
inference_lock = threading.Lock()

def load_model(
    pretrained_model_name_or_path: str = "mutou0308/GSASR",
    model_name: str = "HATL_SA1B",
    device: str | torch.device = "cuda"
):
    enc_path = hf_hub_download(
            repo_id=pretrained_model_name_or_path, filename=os.path.join('GSASR_enhenced_ultra', model_name, 'encoder.pth')
        )
    dec_path = hf_hub_download(
            repo_id=pretrained_model_name_or_path, filename=os.path.join('GSASR_enhenced_ultra', model_name,  'decoder.pth')
        )

    enc_weight = torch.load(enc_path, weights_only=True)['params_ema']
    dec_weight = torch.load(dec_path, weights_only=True)['params_ema']
    
    if model_name in ['EDSR_DIV2K', 'EDSR_DF2K']:
        encoder = EDSRNOUP()
        decoder = Fea2GS_ROPE_AMP()
    elif model_name in ['RDN_DIV2K', 'RDN_DF2K']:
        encoder = RDNNOUP()
        decoder = Fea2GS_ROPE_AMP(num_crossattn_blocks = 2)
    elif model_name in ['SwinIR_DIV2K', 'SwinIR_DF2K']:
        encoder = SwinIRNOUP()
        decoder = Fea2GS_ROPE_AMP(num_crossattn_blocks=2, num_crossattn_layers=4, num_gs_seed=256, window_size=16)
    elif model_name in ['HATL_SA1B']:
        encoder = HATNOUP_ROPE_AMP()
        decoder = Fea2GS_ROPE_AMP(channel=192, num_crossattn_blocks=4, num_crossattn_layers=4, num_selfattn_blocks=8, num_selfattn_layers=6,
                                  num_gs_seed=256, window_size=16)
    else:
        raise ValueError(f"args.model-{model_name} must be in ['EDSR_DIV2K', 'EDSR_DF2K', 'RDN_DIV2K', 'RDN_DF2K', 'SwinIR_DIV2K', 'SwinIR_DF2K', 'HATL_SA1B']")

    encoder.load_state_dict(enc_weight, strict=True)
    decoder.load_state_dict(dec_weight, strict=True)
    encoder.eval()
    decoder.eval()
    encoder = encoder.to(device)
    decoder = decoder.to(device)
    return encoder, decoder


def preprocess(x, denominator=16):
    """Preprocess image to ensure dimensions are multiples of denominator"""
    _, c, h, w = x.shape
    if h % denominator > 0:
        pad_h = denominator - h % denominator
    else:
        pad_h = 0
    if w % denominator > 0:
        pad_w = denominator - w % denominator
    else:
        pad_w = 0
    x_new = F.pad(x, (0, pad_w, 0, pad_h), 'reflect')
    return x_new

def postprocess(x, gt_size_h, gt_size_w):
    """Post-process by cropping to target size"""
    x_new = x[:, :, :gt_size_h, :gt_size_w]
    return x_new

def should_use_tile(image_height, image_width, threshold=1024):
    """Determine if tile processing should be used based on image resolution"""
    return max(image_height, image_width) > threshold

def set_stop_flag():
    """Set the global stop flag to interrupt inference"""
    global stop_inference
    with inference_lock:
        stop_inference = True
    return "🛑 Stopping inference...", gr.update(interactive=False)

def reset_stop_flag():
    """Reset the global stop flag"""
    global stop_inference
    with inference_lock:
        stop_inference = False

def check_stop_flag():
    """Check if inference should be stopped"""
    global stop_inference
    with inference_lock:
        return stop_inference
    
@spaces.GPU
def super_resolution_inference(image, scale=4.0):
    """Super-resolution inference function with automatic tile processing"""

    # Check if gscuda setup has been run
    setup_marker = ".setup_complete"
    if not os.path.exists(setup_marker):
        print("First run detected, installing dependencies...")
        try:
            # subprocess.check_call(["pip", "install", "-e", "."])
            subprocess.check_call(["pip", "install", "dist/gscuda-0.0.0-cp310-cp310-linux_x86_64.whl"])
            # Create marker file to indicate setup is complete
            with open(setup_marker, "w") as f:
                f.write("Setup completed")
            print("Setup completed successfully!")
        except subprocess.CalledProcessError as e:
            return None, f"❌ Setup failed with error: {e}", None


        
    if image is None:
        return None, "Please upload an image", None
    
    # Load model
    encoder, decoder = load_model(model_name="HATL_SA1B")

    # Reset stop flag at the beginning
    reset_stop_flag()
    
    # Fixed parameters
    tile_overlap = 16  # Fixed overlap size
    crop_size = 8     # Fixed crop size
    tile_size = 1024   # Fixed tile size for large images
    
    try:
        # Check for interruption
        if check_stop_flag():
            return None, "❌ Inference interrupted", None
            
        # Convert PIL image to numpy array
        img_np = np.array(image)
        if len(img_np.shape) == 3:
            img_np = img_np[:, :, [2, 1, 0]]  # RGB to BGR
        
        # Convert to tensor
        img = torch.from_numpy(np.transpose(img_np.astype(np.float32) / 255., (2, 0, 1))).float()
        img = img.unsqueeze(0).to(device)
        
        # Check for interruption
        if check_stop_flag():
            return None, "❌ Inference interrupted", None
        
        # Calculate target size
        gt_size = [math.floor(scale * img.shape[2]), math.floor(scale * img.shape[3])]
        
        # Determine if tile processing should be used
        use_tile = should_use_tile(img.shape[2], img.shape[3])
        
        # Force AMP mixed precision
        with torch.inference_mode():
            with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
                # Check for interruption before main processing
                if check_stop_flag():
                    return None, "❌ Inference interrupted", None
                    
                if use_tile:
                    # Use tile processing
                    assert tile_size % 16 == 0, f"tile_size-{tile_size} must be divisible by 16"
                    assert 2 * tile_overlap < tile_size, f"2 * tile_overlap must be less than tile_size"
                    assert 2 * crop_size <= tile_overlap, f"2 * crop_size must be less than or equal to tile_overlap"
                    
                    with torch.no_grad():
                        output = split_and_joint_image(
                            lq=img, 
                            scale_factor=scale,
                            split_size=tile_size,
                            overlap_size=tile_overlap,
                            model_g=encoder,
                            model_fea2gs=decoder,
                            crop_size=crop_size,
                            scale_modify=torch.tensor([scale, scale]),
                            default_step_size=1.2,
                            cuda_rendering=True,
                            mode='scale_modify',
                            if_dmax=True,
                            dmax_mode='fix',
                            dmax=0.1
                        )
                else:
                    # Direct processing without tiles
                    lq_pad = preprocess(img, 16)  # denominator=16 for HATL
                    gt_size_pad = torch.tensor([math.floor(scale * lq_pad.shape[2]), 
                                            math.floor(scale * lq_pad.shape[3])])
                    gt_size_pad = gt_size_pad.unsqueeze(0)
                    
                    with torch.no_grad():
                        # Check for interruption before encoder
                        if check_stop_flag():
                            return None, "❌ Inference interrupted", None
                            
                        # Encoder output
                        encoder_output = encoder(lq_pad)  # b,c,h,w
                        
                        # Check for interruption before decoder
                        if check_stop_flag():
                            return None, "❌ Inference interrupted", None
                            
                        scale_vector = torch.tensor(scale, dtype=torch.float32).unsqueeze(0).to(device)
                        
                        # Decoder output
                        batch_gs_parameters = decoder(encoder_output, scale_vector)
                        gs_parameters = batch_gs_parameters[0, :]
                        
                        # Check for interruption before gaussian rendering
                        if check_stop_flag():
                            return None, "❌ Inference interrupted", None
                        
                        # Gaussian rendering
                        b_output = generate_2D_gaussian_splatting_step(
                            gs_parameters=gs_parameters,
                            sr_size=gt_size_pad[0],
                            scale=scale,
                            sample_coords=None,
                            scale_modify=torch.tensor([scale, scale]),
                            default_step_size=1.2,
                            cuda_rendering=True,
                            mode='scale_modify',
                            if_dmax=True,
                            dmax_mode='fix',
                            dmax=0.1
                        )
                        output = b_output.unsqueeze(0)
        
        # Check for interruption before post-processing
        if check_stop_flag():
            return None, "❌ Inference interrupted", None
        
        # Post-processing
        output = postprocess(output, gt_size[0], gt_size[1])
        
        # Convert back to PIL image format
        output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
        output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))  # BGR to RGB
        output = (output * 255.0).round().astype(np.uint8)
        
        # Convert to PIL image
        output_pil = Image.fromarray(output)
        
        # Generate result information
        original_size = f"{img.shape[3]}x{img.shape[2]}"
        output_size = f"{output.shape[1]}x{output.shape[0]}"
        tile_info = f"Tile processing enabled (size: {tile_size})" if use_tile else "Direct processing (no tiles)"
        result_info = f"✅ Processing completed successfully!\nOriginal size: {original_size}\nSuper-resolution size: {output_size}\nScale factor: {scale:.2f}x\nProcessing mode: {tile_info}\nAMP acceleration: Force enabled\nOverlap size: {tile_overlap}\nCrop size: {crop_size}"
        
        return output_pil, result_info, output_pil
        
    except Exception as e:
        if check_stop_flag():
            return None, "❌ Inference interrupted", None
        return None, f"❌ Error during processing: {str(e)}", None

def predict(image, scale):
    """Gradio prediction function"""
    output_image, info, download_image = super_resolution_inference(image, scale)
    
    # If processing successful, save image for download
    if output_image is not None:
        # Create temporary filename
        timestamp = int(time.time())
        temp_filename = f"GSASR_SR_result_{scale}x_{timestamp}.png"
        temp_path = os.path.join(tempfile.gettempdir(), temp_filename)
        
        # Save image
        output_image.save(temp_path, "PNG")
        
        return output_image, temp_path, "✅ Ready", gr.update(interactive=True)
    else:
        return output_image, None, info if info else "❌ Processing failed", gr.update(interactive=True)

# Create Gradio interface
with gr.Blocks(title="🚀 GSASR (2D Gaussian Splatting Super-Resolution)") as demo:
    gr.Markdown("# **🚀 GSASR (Generalized and efficient 2d gaussian splatting for arbitrary-scale super-resolution)**")
    gr.Markdown("Official demo for GSASR. Please refer to our [paper](https://arxiv.org/pdf/2501.06838), [project page](https://mt-cly.github.io/GSASR.github.io/), and [github](https://github.com/ChrisDud0257/GSASR) for more details.")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Image")
            
            # Scale parameters
            with gr.Group():
                gr.Markdown("### SR Scale")
                scale_slider = gr.Slider(minimum=1.0, maximum=30.0, value=4.0, step=0.1, label="SR Scale")
            
            # Control buttons
            with gr.Row():
                submit_btn = gr.Button("🚀 Start Super-Resolution", variant="primary")
                stop_btn = gr.Button("🛑 Stop Inference", variant="stop")
        
        with gr.Column():
            output_image = gr.Image(type="pil", label="Super-Resolution Result")
            
            # Status display
            status_text = gr.Textbox(label="Status", value="✅ Ready", interactive=False)
            
            # Download component
            with gr.Group():
                gr.Markdown("### 📥 Download Super-Resolution Result")
                download_btn = gr.File(visible=True)
    
    # Event handlers
    submit_event = submit_btn.click(
        fn=predict,
        inputs=[input_image, scale_slider],
        outputs=[output_image, download_btn, status_text, stop_btn]
    )
    
    stop_btn.click(
        fn=set_stop_flag,
        inputs=[],
        outputs=[status_text, stop_btn],
        cancels=[submit_event]
    )
    
    # Example images
    gr.Markdown("### 📚 Example Images")
    gr.Markdown("Try these examples with different scales:")
    
    gr.Examples(
        examples=[
            ["assets/0846x4.png", 1.5],
            ["assets/0892x4.png", 2.8],
            ["assets/0873x4_cropped_120x120.png", 30.0]
        ],
        inputs=[input_image, scale_slider],
        examples_per_page=3,
        cache_examples=False,
        label="Examples"
    )

if __name__ == "__main__":
    demo.launch(share=True, server_name="0.0.0.0")