Update app.py

app.py

@@ -49,14 +49,17 @@ def initialize_models():
         # Load pipeline
         pipeline = load_pipeline(args, accelerator, enable_xformers_memory_efficient_attention=False)
         
-        # Set pipeline to eval mode
+        # 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(device=accelerator.device)
+        generator = torch.Generator("cuda")
         
         return True
 
@@ -64,7 +67,7 @@ def initialize_models():
         print(f"Error initializing models: {str(e)}")
         return False
 
-@spaces.GPU(processing_timeout=180) # Increased timeout for longer processing
+@spaces.GPU(processing_timeout=180)
 def process_image(
     input_image,
     prompt="clean, high-resolution, 8k",
@@ -78,11 +81,12 @@ def process_image(
 ):
     global pipeline, generator, accelerator
     
-    if pipeline is None:
-        if not initialize_models():
-            return None
-
     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,
@@ -124,27 +128,38 @@ 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=args.use_vae_encode_condition,
-            )
+            try:
+                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,
+                )
+            except Exception as e:
+                print(f"Pipeline execution error: {str(e)}")
+                raise
 
         image = output.images[0]
 
@@ -161,6 +176,8 @@ def process_image(
 
     except Exception as e:
         print(f"Error processing image: {str(e)}")
+        import traceback
+        traceback.print_exc()
         return None
 
 # Create Gradio interface