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)

@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
        pipeline.unet.eval()
        pipeline.controlnet.eval()
        pipeline.vae.eval()
        pipeline.text_encoder.eval()
        
        # Move pipeline to CUDA
        pipeline = pipeline.to("cuda")
        
        # Initialize generator
        generator = torch.Generator("cuda")
        
        return True

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

@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, accelerator
    
    try:
        # Initialize models if not already done
        if pipeline is None:
            if not initialize_models():
                return None

        # 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,  # Changed to True
            sample_times=1
        )

        # Set seed if provided
        if seed is not None:
            generator.manual_seed(seed)

        # 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

        # Ensure pipeline is on CUDA and in eval mode
        pipeline = pipeline.to("cuda")
        pipeline.unet.eval()
        pipeline.controlnet.eval()
        pipeline.vae.eval()
        pipeline.text_encoder.eval()

        # Generate image
        with torch.no_grad():
            try:
                # First encode the image with VAE
                image_tensor = pipeline.image_processor.preprocess(validation_image)
                image_tensor = image_tensor.unsqueeze(0).to(device="cuda", dtype=torch.float32)
                
                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,  # Set to True
                )
            except Exception as e:
                print(f"Pipeline execution error: {str(e)}")
                raise

        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": 2.5,
    "conditioning_scale": 1.0,
    "num_steps": 6,
    "seed": None,
    "upscale_factor": 4,
    "color_fix_method": "adain"
}

# Create interface components
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input Image")
            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=4, value=DEFAULT_VALUES["upscale_factor"], step=1, label="Upscale Factor")
            color_fix_method = gr.Radio(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()