Update app.py

app.py

@@ -18,6 +18,7 @@ class Args:
     def __init__(self, **kwargs):
         self.__dict__.update(kwargs)
 
+# Initialize models at startup
 @spaces.GPU
 def initialize_models():
     global pipeline, generator, accelerator
@@ -49,24 +50,27 @@ def initialize_models():
         # Load pipeline
         pipeline = load_pipeline(args, accelerator, enable_xformers_memory_efficient_attention=False)
         
-        # Ensure all models are in eval mode
+        # 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()
         
-        # Move pipeline to CUDA
-        pipeline = pipeline.to("cuda")
-        
         # 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,
@@ -79,14 +83,9 @@ def process_image(
     upscale_factor=4,
     color_fix_method="adain"
 ):
-    global pipeline, generator, accelerator
+    global pipeline, generator
     
     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,
@@ -105,7 +104,7 @@ 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
         )
 
@@ -128,42 +127,27 @@ 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 = 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]
 
@@ -184,7 +168,6 @@ def process_image(
         traceback.print_exc()
         return None
 
-
 # Define default values
 DEFAULT_VALUES = {
     "prompt": "clean, texture, high-resolution, 8k",