app.py

import os
import torch
import gradio as gr
import spaces
from PIL import Image
from diffusers import DiffusionPipeline
from huggingface_hub import snapshot_download
from test_ccsr_tile import load_pipeline
import argparse
from accelerate import Accelerator

# Initialize global variables
pipeline = None
generator = None
accelerator = None

class Args:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)

# Initialize models at startup
@spaces.GPU
def initialize_models():
    global pipeline, generator, accelerator
    
    try:
        # Download model repository
        model_path = snapshot_download(
            repo_id="NightRaven109/CCSRModels",
            token=os.environ['Read2']
        )

        # Set up default arguments
        args = Args(
            pretrained_model_path=os.path.join(model_path, "stable-diffusion-2-1-base"),
            controlnet_model_path=os.path.join(model_path, "Controlnet"),
            vae_model_path=os.path.join(model_path, "vae"),
            mixed_precision="fp16",
            tile_vae=False,
            sample_method="ddpm",
            vae_encoder_tile_size=1024,
            vae_decoder_tile_size=224
        )

        # Initialize accelerator
        accelerator = Accelerator(
            mixed_precision=args.mixed_precision,
        )

        # Load pipeline
        pipeline = load_pipeline(args, accelerator, enable_xformers_memory_efficient_attention=False)
        
        # Ensure all models are in eval mode and on CUDA
        pipeline = pipeline.to("cuda")
        pipeline.unet.eval()
        pipeline.controlnet.eval()
        pipeline.vae.eval()
        pipeline.text_encoder.eval()
        
        # Initialize generator
        generator = torch.Generator("cuda")
        
        print("Models initialized and ready!")
        return True

    except Exception as e:
        print(f"Error initializing models: {str(e)}")
        return False

# Load models at module level
print("Initializing models...")
initialize_models()

@spaces.GPU(processing_timeout=180)
def process_image(
    input_image,
    prompt="clean, texture, high-resolution, 8k",
    negative_prompt="blurry, dotted, noise, raster lines, unclear, lowres, over-smoothed",
    guidance_scale=2.5,
    conditioning_scale=1.0,
    num_inference_steps=6,
    seed=None,
    upscale_factor=4,
    color_fix_method="adain"
):
    global pipeline, generator
    
    try:
        # Handle seed
        if seed is not None and seed != 0:  # Only set seed if it's provided and not 0
            if generator is None:
                generator = torch.Generator("cuda")
            generator.manual_seed(seed)
        elif generator is None:
            generator = torch.Generator("cuda")

        # Create args object with all necessary parameters
        args = Args(
            added_prompt=prompt,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            conditioning_scale=conditioning_scale,
            num_inference_steps=num_inference_steps,
            seed=seed,
            upscale=upscale_factor,
            process_size=512,
            align_method=color_fix_method,
            t_max=0.6666,
            t_min=0.0,
            tile_diffusion=False,
            tile_diffusion_size=None,
            tile_diffusion_stride=None,
            start_steps=999,
            start_point='lr',
            use_vae_encode_condition=True,
            sample_times=1
        )

        # Process input image
        validation_image = Image.fromarray(input_image)
        ori_width, ori_height = validation_image.size
        
        # Resize logic
        resize_flag = False
        if ori_width < args.process_size//args.upscale or ori_height < args.process_size//args.upscale:
            scale = (args.process_size//args.upscale)/min(ori_width, ori_height)
            validation_image = validation_image.resize((round(scale*ori_width), round(scale*ori_height)))
            resize_flag = True

        validation_image = validation_image.resize((validation_image.size[0]*args.upscale, validation_image.size[1]*args.upscale))
        validation_image = validation_image.resize((validation_image.size[0]//8*8, validation_image.size[1]//8*8))
        width, height = validation_image.size

        # Generate image
        with torch.no_grad():
            inference_time, output = pipeline(
                args.t_max,
                args.t_min,
                args.tile_diffusion,
                args.tile_diffusion_size,
                args.tile_diffusion_stride,
                args.added_prompt,
                validation_image,
                num_inference_steps=args.num_inference_steps,
                generator=generator,
                height=height,
                width=width,
                guidance_scale=args.guidance_scale,
                negative_prompt=args.negative_prompt,
                conditioning_scale=args.conditioning_scale,
                start_steps=args.start_steps,
                start_point=args.start_point,
                use_vae_encode_condition=True,
            )

        image = output.images[0]

        # Apply color fixing if specified
        if args.align_method != "none":
            from myutils.wavelet_color_fix import wavelet_color_fix, adain_color_fix
            fix_func = wavelet_color_fix if args.align_method == "wavelet" else adain_color_fix
            image = fix_func(image, validation_image)
            
        if resize_flag:
            image = image.resize((ori_width*args.upscale, ori_height*args.upscale))

        return image

    except Exception as e:
        print(f"Error processing image: {str(e)}")
        import traceback
        traceback.print_exc()
        return None

# Define default values
DEFAULT_VALUES = {
    "prompt": "clean, texture, high-resolution, 8k",
    "negative_prompt": "blurry, dotted, noise, raster lines, unclear, lowres, over-smoothed",
    "guidance_scale": 3,
    "conditioning_scale": 1.0,
    "num_steps": 6,
    "seed": None,
    "upscale_factor": 4,
    "color_fix_method": "adain"
}

# Create interface components
with gr.Blocks(title="Controllable Conditional Super-Resolution") as demo:
    gr.Markdown("## Controllable Conditional Super-Resolution")
    gr.Markdown("Upload an image to enhance its resolution using CCSR.")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input Image")
            
            # Put all parameters in an accordion/dropdown
            with gr.Accordion("Advanced Options", open=False):
                prompt = gr.Textbox(label="Prompt", value=DEFAULT_VALUES["prompt"])
                negative_prompt = gr.Textbox(label="Negative Prompt", value=DEFAULT_VALUES["negative_prompt"])
                guidance_scale = gr.Slider(minimum=1.0, maximum=20.0, value=DEFAULT_VALUES["guidance_scale"], label="Guidance Scale")
                conditioning_scale = gr.Slider(minimum=0.1, maximum=2.0, value=DEFAULT_VALUES["conditioning_scale"], label="Conditioning Scale")
                num_steps = gr.Slider(minimum=1, maximum=50, value=DEFAULT_VALUES["num_steps"], step=1, label="Number of Steps")
                seed = gr.Number(label="Seed", value=DEFAULT_VALUES["seed"])
                upscale_factor = gr.Slider(minimum=1, maximum=8, value=DEFAULT_VALUES["upscale_factor"], step=1, label="Upscale Factor")
                color_fix_method = gr.Dropdown(
                    choices=["none", "wavelet", "adain"], 
                    label="Color Fix Method", 
                    value=DEFAULT_VALUES["color_fix_method"]
                )
            
            # Add buttons
            with gr.Row():
                clear_btn = gr.Button("Clear")
                submit_btn = gr.Button("Submit", variant="primary")

        with gr.Column():
            output_image = gr.Image(label="Generated Image")

    # Define submit action
    submit_btn.click(
        fn=process_image,
        inputs=[
            input_image, prompt, negative_prompt, guidance_scale,
            conditioning_scale, num_steps, seed, upscale_factor,
            color_fix_method
        ],
        outputs=output_image
    )

    # Define clear action that resets to default values
    def reset_to_defaults():
        return [
            None,  # input_image
            DEFAULT_VALUES["prompt"],
            DEFAULT_VALUES["negative_prompt"],
            DEFAULT_VALUES["guidance_scale"],
            DEFAULT_VALUES["conditioning_scale"],
            DEFAULT_VALUES["num_steps"],
            DEFAULT_VALUES["seed"],
            DEFAULT_VALUES["upscale_factor"],
            DEFAULT_VALUES["color_fix_method"]
        ]

    clear_btn.click(
        fn=reset_to_defaults,
        inputs=None,
        outputs=[
            input_image, prompt, negative_prompt, guidance_scale,
            conditioning_scale, num_steps, seed, upscale_factor,
            color_fix_method
        ]
    )

if __name__ == "__main__":
    demo.launch()