Update app.py

app.py

@@ -3,8 +3,9 @@ import torch
 import gradio as gr
 import spaces
 from PIL import Image
+from diffusers import DiffusionPipeline
 from huggingface_hub import snapshot_download
-from test_ccsr_tile import main, load_pipeline
+from test_ccsr_tile import load_pipeline
 import argparse
 from accelerate import Accelerator
 
@@ -48,6 +49,12 @@ def initialize_models():
         # 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()
+        
         # Initialize generator
         generator = torch.Generator(device=accelerator.device)
         
@@ -57,7 +64,7 @@ def initialize_models():
         print(f"Error initializing models: {str(e)}")
         return False
 
-@spaces.GPU
+@spaces.GPU(processing_timeout=180) # Increased timeout for longer processing
 def process_image(
     input_image,
     prompt="clean, high-resolution, 8k",
@@ -117,13 +124,6 @@ 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
         with torch.no_grad():
             inference_time, output = pipeline(
@@ -157,62 +157,30 @@ 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
 
-# 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
+# Create Gradio interface
+demo = 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__":
+    demo.launch()