Update app.py

app.py

@@ -8,7 +8,11 @@ import safetensors.torch as sf
 from PIL import Image
 from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
 from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler, EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler
-from diffusers.models.attention_processor import AttnProcessor2_0
+try:
+    from diffusers.models.attention_processor import AttnProcessor2_0
+except ImportError:
+    # Fallback for older diffusers versions
+    AttnProcessor2_0 = None
 from transformers import CLIPTextModel, CLIPTokenizer
 from enum import Enum
 from torch.hub import download_url_to_file
@@ -20,6 +24,7 @@ from huggingface_hub import PyTorchModelHubMixin
 try:
     from transformers import pipeline
     rmbg_pipeline = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
+    USE_RMBG_PIPELINE = True
 except Exception as e:
     print(f"Failed to load RMBG pipeline: {e}")
     USE_RMBG_PIPELINE = False
@@ -46,10 +51,10 @@ print(f"Using device: {device}")
 print("Loading models...")
 
 # Initialize models
-tokenizer = CLIPTokenizer.from_pretrained(sd15_name, subfolder="tokenizer", trust_remote_code=True)
-text_encoder = CLIPTextModel.from_pretrained(sd15_name, subfolder="text_encoder", trust_remote_code=True)
-vae = AutoencoderKL.from_pretrained(sd15_name, subfolder="vae", trust_remote_code=True)
-unet = UNet2DConditionModel.from_pretrained(sd15_name, subfolder="unet", trust_remote_code=True)
+tokenizer = CLIPTokenizer.from_pretrained(sd15_name, subfolder="tokenizer")
+text_encoder = CLIPTextModel.from_pretrained(sd15_name, subfolder="text_encoder")
+vae = AutoencoderKL.from_pretrained(sd15_name, subfolder="vae")
+unet = UNet2DConditionModel.from_pretrained(sd15_name, subfolder="unet")
 
 # Modify UNet for IC-Light
 with torch.no_grad():
@@ -88,8 +93,14 @@ vae = vae.to(device=device, dtype=torch.bfloat16)
 unet = unet.to(device=device, dtype=torch.float16)
 
 # Set attention processors
-unet.set_attn_processor(AttnProcessor2_0())
-vae.set_attn_processor(AttnProcessor2_0())
+if AttnProcessor2_0 is not None:
+    try:
+        unet.set_attn_processor(AttnProcessor2_0())
+        vae.set_attn_processor(AttnProcessor2_0())
+    except Exception as e:
+        print(f"Failed to set attention processors: {e}")
+else:
+    print("AttnProcessor2_0 not available, using default processors")
 
 # Scheduler
 scheduler = DPMSolverMultistepScheduler(
@@ -340,10 +351,66 @@ def process(input_fg, input_bg, prompt, image_width, image_height, num_samples,
 
 @torch.inference_mode()
 def process_relight(input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source):
-    input_fg, matting = run_rmbg(input_fg)
-    results, extra_images = process(input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source)
-    results = [(x * 255.0).clip(0, 255).astype(np.uint8) for x in results]
-    return results + extra_images
+    try:
+        # Input validation
+        if input_fg is None:
+            raise ValueError("Please upload a foreground image")
+        if input_bg is None and bg_source == "Use Background Image":
+            raise ValueError("Please upload a background image or choose a lighting direction")
+        if not prompt.strip():
+            raise ValueError("Please enter a prompt")
+        
+        print(f"Processing with device: {device}")
+        print(f"Input shapes - FG: {input_fg.shape}, BG: {input_bg.shape if input_bg is not None else 'None'}")
+        
+        # Optimize for Hugging Face free GPU (limited memory)
+        if device.type == 'cuda':
+            # Limit image size for free GPU tier
+            max_size = 768  # Increased for GPU but still conservative
+            if image_width > max_size or image_height > max_size:
+                scale = min(max_size / image_width, max_size / image_height)
+                image_width = int(image_width * scale // 64) * 64  # Keep multiple of 64
+                image_height = int(image_height * scale // 64) * 64
+                print(f"Reduced image size for GPU memory: {image_width}x{image_height}")
+            
+            # Disable highres for free tier to save memory
+            if highres_scale > 1.0:
+                highres_scale = 1.0
+                print("Disabled highres scaling to save GPU memory")
+                
+        elif device.type == 'cpu':
+            # Limit image size for CPU processing
+            max_size = 512
+            if image_width > max_size or image_height > max_size:
+                image_width = min(image_width, max_size)
+                image_height = min(image_height, max_size)
+                print(f"Reduced image size for CPU: {image_width}x{image_height}")
+            
+            # Limit number of samples for CPU
+            if num_samples > 1:
+                num_samples = 1
+                print("Reduced num_samples to 1 for CPU processing")
+        
+        print("Running background removal...")
+        input_fg, matting = run_rmbg(input_fg)
+        
+        print("Starting main processing...")
+        results, extra_images = process(input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source)
+        
+        print("Converting results...")
+        results = [(x * 255.0).clip(0, 255).astype(np.uint8) for x in results]
+        
+        print("Processing completed successfully!")
+        return results + extra_images
+        
+    except Exception as e:
+        print(f"Error in process_relight: {str(e)}")
+        import traceback
+        traceback.print_exc()
+        # Return error image
+        error_img = np.zeros((512, 512, 3), dtype=np.uint8)
+        error_img[:, :] = [255, 0, 0]  # Red error image
+        return [error_img]
 
 # Quick prompts for easy testing
 quick_prompts = [
@@ -405,10 +472,16 @@ def create_demo():
         
         # Event handlers
         inputs = [input_fg, input_bg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, bg_source]
-        relight_button.click(fn=process_relight, inputs=inputs, outputs=[result_gallery])
+        relight_button.click(
+            fn=process_relight, 
+            inputs=inputs, 
+            outputs=[result_gallery],
+            show_progress=True,
+            queue=True
+        )
         example_prompts.click(lambda x: x[0], inputs=example_prompts, outputs=prompt, show_progress=False, queue=False)
         
-        # Examples
+        # Examples - temporarily disabled due to missing image files
         # gr.Examples(
         #     examples=[
         #         ["examples/person1.jpg", "examples/bg1.jpg", "beautiful woman, cinematic lighting", "Use Background Image"],