Update app.py

app.py

@@ -9,10 +9,15 @@ from test_ccsr_tile import load_pipeline
 import argparse
 from accelerate import Accelerator
 
-# Initialize global variables
-pipeline = None
-generator = None
-accelerator = None
+# Global variables
+class ModelContainer:
+    def __init__(self):
+        self.pipeline = None
+        self.generator = None
+        self.accelerator = None
+        self.is_initialized = False
+
+model_container = ModelContainer()
 
 class Args:
     def __init__(self, **kwargs):
@@ -20,10 +25,12 @@ class Args:
 
 @spaces.GPU
 def initialize_models():
-    global pipeline, generator, accelerator
+    """Initialize models only if they haven't been initialized yet"""
+    if model_container.is_initialized:
+        return True
     
     try:
-        # Download model repository
+        # Download model repository (only once)
         model_path = snapshot_download(
             repo_id="NightRaven109/CCSRModels",
             token=os.environ['Read2']
@@ -42,24 +49,28 @@ def initialize_models():
         )
 
         # Initialize accelerator
-        accelerator = Accelerator(
+        model_container.accelerator = Accelerator(
             mixed_precision=args.mixed_precision,
         )
 
         # Load pipeline
-        pipeline = load_pipeline(args, accelerator, enable_xformers_memory_efficient_attention=False)
+        model_container.pipeline = load_pipeline(args, model_container.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()
+        # Set models to eval mode
+        model_container.pipeline.unet.eval()
+        model_container.pipeline.controlnet.eval()
+        model_container.pipeline.vae.eval()
+        model_container.pipeline.text_encoder.eval()
         
-        # Move pipeline to CUDA
-        pipeline = pipeline.to("cuda")
+        # Move pipeline to CUDA and set to eval mode once
+        model_container.pipeline = model_container.pipeline.to("cuda")
         
         # Initialize generator
-        generator = torch.Generator("cuda")
+        model_container.generator = torch.Generator("cuda")
+        
+        # Set initialization flag
+        model_container.is_initialized = True
         
         return True
 
@@ -67,6 +78,7 @@ def initialize_models():
         print(f"Error initializing models: {str(e)}")
         return False
 
+@torch.no_grad()  # Add no_grad decorator for inference
 @spaces.GPU(processing_timeout=180)
 def process_image(
     input_image,
@@ -79,15 +91,13 @@ def process_image(
     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
+    # Initialize models if not already done
+    if not model_container.is_initialized:
+        if not initialize_models():
+            return None
 
-        # Create args object with all necessary parameters
+    try:
+        # Create args object
         args = Args(
             added_prompt=prompt,
             negative_prompt=negative_prompt,
@@ -105,13 +115,13 @@ def process_image(
             tile_diffusion_stride=None,
             start_steps=999,
             start_point='lr',
-            use_vae_encode_condition=True,  # Changed to True
+            use_vae_encode_condition=True,
             sample_times=1
         )
 
         # Set seed if provided
         if seed is not None:
-            generator.manual_seed(seed)
+            model_container.generator.manual_seed(seed)
 
         # Process input image
         validation_image = Image.fromarray(input_image)
@@ -128,42 +138,26 @@ def process_image(
         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
+        inference_time, output = model_container.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=model_container.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]