Update app.py

app.py

@@ -2,115 +2,52 @@ import os
 import torch
 import gradio as gr
 import spaces
-import numpy as np
 from PIL import Image
-import safetensors.torch
 from huggingface_hub import snapshot_download
+from test_ccsr_tile import main, load_pipeline
+import argparse
 from accelerate import Accelerator
-from accelerate.utils import set_seed
-from diffusers import (
-    AutoencoderKL,
-    DDPMScheduler,
-    UNet2DConditionModel,
-)
-from transformers import CLIPTextModel, CLIPTokenizer, CLIPImageProcessor
-from models.controlnet import ControlNetModel
-from pipelines.pipeline_ccsr import StableDiffusionControlNetPipeline
-from myutils.wavelet_color_fix import wavelet_color_fix, adain_color_fix
 
-# Initialize global variables for models
+# Initialize global variables
 pipeline = None
 generator = None
 accelerator = None
-model_path = None
-
-def load_pipeline(accelerator, model_path):
-    # Load scheduler
-    scheduler = DDPMScheduler.from_pretrained(
-        model_path,
-        subfolder="stable-diffusion-2-1-base/scheduler"
-    )
-
-    # Load models
-    text_encoder = CLIPTextModel.from_pretrained(
-        model_path,
-        subfolder="stable-diffusion-2-1-base/text_encoder"
-    )
-    
-    tokenizer = CLIPTokenizer.from_pretrained(
-        model_path,
-        subfolder="stable-diffusion-2-1-base/tokenizer"
-    )
-    
-    feature_extractor = CLIPImageProcessor.from_pretrained(
-        os.path.join(model_path, "stable-diffusion-2-1-base/feature_extractor")
-    )
-    
-    unet = UNet2DConditionModel.from_pretrained(
-        model_path,
-        subfolder="stable-diffusion-2-1-base/unet"
-    )
-    
-    controlnet = ControlNetModel.from_pretrained(
-        model_path,
-        subfolder="Controlnet"
-    )
-    
-    vae = AutoencoderKL.from_pretrained(
-        model_path,
-        subfolder="vae"
-    )
-
-    # Freeze models
-    for model in [vae, text_encoder, unet, controlnet]:
-        model.requires_grad_(False)
-
-    # Initialize pipeline
-    pipeline = StableDiffusionControlNetPipeline(
-        vae=vae,
-        text_encoder=text_encoder,
-        tokenizer=tokenizer,
-        feature_extractor=feature_extractor,
-        unet=unet,
-        controlnet=controlnet,
-        scheduler=scheduler,
-        safety_checker=None,
-        requires_safety_checker=False,
-    )
-
-    # Set weight dtype based on mixed precision
-    weight_dtype = torch.float32
-    if accelerator.mixed_precision == "fp16":
-        weight_dtype = torch.float16
-    elif accelerator.mixed_precision == "bf16":
-        weight_dtype = torch.bfloat16
-
-    # Move models to accelerator device with appropriate dtype
-    for model in [text_encoder, vae, unet, controlnet]:
-        model.to(accelerator.device, dtype=weight_dtype)
-
-    return pipeline
+
+class Args:
+    def __init__(self, **kwargs):
+        self.__dict__.update(kwargs)
 
 @spaces.GPU
 def initialize_models():
-    global pipeline, generator, accelerator, model_path
+    global pipeline, generator, accelerator
     
-    # Initialize accelerator
-    accelerator = Accelerator(
-        mixed_precision="fp16",
-        gradient_accumulation_steps=1
-    )
-
     try:
-        # Download the entire repository
+        # Download model repository
         model_path = snapshot_download(
             repo_id="NightRaven109/CCSRModels",
             token=os.environ['Read2']
         )
-        
-        # Load pipeline using the original loading function
-        pipeline = load_pipeline(accelerator, model_path)
 
+        # 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)
+        
         # Initialize generator
         generator = torch.Generator(device=accelerator.device)
         
@@ -137,72 +74,83 @@ def process_image(
     if pipeline is None:
         if not initialize_models():
             return None
-    
+
     try:
-        # Set seed
+        # 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=False,
+            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 from original script
+        # Resize logic
         resize_flag = False
-        rscale = upscale_factor
-        process_size = 512  # Same as args.process_size in original
-        
-        if ori_width < process_size//rscale or ori_height < process_size//rscale:
-            scale = (process_size//rscale)/min(ori_width, ori_height)
-            tmp_image = validation_image.resize((round(scale*ori_width), round(scale*ori_height)))
-            validation_image = tmp_image
+        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]*rscale, validation_image.size[1]*rscale))
+        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
-        
-        # Move pipeline to GPU for processing
-        pipeline.to(accelerator.device)
-        
+
         # Generate image
-        with torch.no_grad():
-            inference_time, output = pipeline(
-                0.6666,  # t_max
-                0.0,    # t_min
-                False,  # tile_diffusion
-                None,   # tile_diffusion_size
-                None,   # tile_diffusion_stride
-                prompt,
-                validation_image,
-                num_inference_steps=num_inference_steps,
-                generator=generator,
-                height=height,
-                width=width,
-                guidance_scale=guidance_scale,
-                negative_prompt=negative_prompt,
-                conditioning_scale=conditioning_scale,
-                start_steps=999,
-                start_point='lr',
-                use_vae_encode_condition=False
-            )
-        
+        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=args.use_vae_encode_condition,
+        )
+
         image = output.images[0]
 
         # Apply color fixing if specified
-        if color_fix_method != "none":
-            fix_func = wavelet_color_fix if color_fix_method == "wavelet" else adain_color_fix
+        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*rscale, ori_height*rscale))
-        
-        # Move pipeline back to CPU
-        pipeline.to("cpu")
-        torch.cuda.empty_cache()
-        
+            image = image.resize((ori_width*args.upscale, ori_height*args.upscale))
+
         return image
-    
+
     except Exception as e:
         print(f"Error processing image: {str(e)}")
         return None