Update app.py

app.py

@@ -18,7 +18,6 @@ class Args:
     def __init__(self, **kwargs):
         self.__dict__.update(kwargs)
 
-# Initialize models at startup
 @spaces.GPU
 def initialize_models():
     global pipeline, generator, accelerator
@@ -50,27 +49,24 @@ def initialize_models():
         # 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")
+        # 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")
         
-        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,
@@ -83,16 +79,13 @@ def process_image(
     upscale_factor=4,
     color_fix_method="adain"
 ):
-    global pipeline, generator
+    global pipeline, generator, accelerator
     
     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")
+        # 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(
@@ -112,10 +105,14 @@ def process_image(
             tile_diffusion_stride=None,
             start_steps=999,
             start_point='lr',
-            use_vae_encode_condition=True,
+            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
@@ -131,27 +128,42 @@ 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():
-            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,
-            )
+            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]
 
@@ -172,6 +184,7 @@ def process_image(
         traceback.print_exc()
         return None
 
+
 # Define default values
 DEFAULT_VALUES = {
     "prompt": "clean, texture, high-resolution, 8k",