Update app.py

app.py

@@ -117,26 +117,34 @@ 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
 
+        # Move pipeline to GPU and set to eval mode
+        pipeline.to(accelerator.device)
+        pipeline.unet.eval()
+        pipeline.controlnet.eval()
+        pipeline.vae.eval()
+        pipeline.text_encoder.eval()
+
         # Generate image
-        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,
-        )
+        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,
+            )
 
         image = output.images[0]
 
@@ -149,30 +157,62 @@ def process_image(
         if resize_flag:
             image = image.resize((ori_width*args.upscale, ori_height*args.upscale))
 
+        # Move pipeline back to CPU to free up GPU memory
+        pipeline.to("cpu")
+        torch.cuda.empty_cache()
+
         return image
 
     except Exception as e:
         print(f"Error processing image: {str(e)}")
         return None
 
-# Create Gradio interface
-iface = gr.Interface(
-    fn=process_image,
-    inputs=[
-        gr.Image(label="Input Image"),
-        gr.Textbox(label="Prompt", value="clean, high-resolution, 8k"),
-        gr.Textbox(label="Negative Prompt", value="blurry, dotted, noise, raster lines, unclear, lowres, over-smoothed"),
-        gr.Slider(minimum=1.0, maximum=20.0, value=1.0, label="Guidance Scale"),
-        gr.Slider(minimum=0.1, maximum=2.0, value=1.0, label="Conditioning Scale"),
-        gr.Slider(minimum=1, maximum=50, value=20, step=1, label="Number of Steps"),
-        gr.Number(label="Seed", value=42),
-        gr.Slider(minimum=1, maximum=4, value=2, step=1, label="Upscale Factor"),
-        gr.Radio(["none", "wavelet", "adain"], label="Color Fix Method", value="adain"),
-    ],
-    outputs=gr.Image(label="Generated Image"),
-    title="Controllable Conditional Super-Resolution",
-    description="Upload an image to enhance its resolution using CCSR."
-)
-
-if __name__ == "__main__":
-    iface.launch()
+# Also update the initialize_models function:
+@spaces.GPU
+def initialize_models():
+    global pipeline, generator, accelerator
+    
+    try:
+        # Download model repository
+        model_path = snapshot_download(
+            repo_id="NightRaven109/CCSRModels",
+            token=os.environ['Read2']
+        )
+
+        # 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)
+        
+        # Set pipeline to eval mode
+        pipeline.unet.eval()
+        pipeline.controlnet.eval()
+        pipeline.vae.eval()
+        pipeline.text_encoder.eval()
+        
+        # Move to CPU initially to save memory
+        pipeline.to("cpu")
+        
+        # Initialize generator
+        generator = torch.Generator(device=accelerator.device)
+        
+        return True
+
+    except Exception as e:
+        print(f"Error initializing models: {str(e)}")
+        return False