Update app.py

app.py

@@ -372,13 +372,108 @@ def generate_ai_background(
     seed: Optional[int] = None,
 ) -> str:
     """Generate AI background using Stable Diffusion."""
-    # TEMPORARILY DISABLED due to PyTorch/Diffusers compatibility issue
-    # To fix: pip install --upgrade torch diffusers transformers
-    raise RuntimeError(
-        "AI Background temporarily disabled due to PyTorch/Diffusers version compatibility.\n"
-        "To fix: pip install --upgrade torch diffusers transformers accelerate\n"
-        "For now, please use Upload Image or Gradients instead."
-    )
+    if not TORCH_AVAILABLE:
+        raise RuntimeError("PyTorch required for AI background generation")
+    
+    try:
+        from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
+    except ImportError as e:
+        raise RuntimeError(f"Please install diffusers: pip install diffusers transformers accelerate\nError: {e}")
+    
+    device = "cuda" if CUDA_AVAILABLE else "cpu"
+    torch_dtype = torch.float16 if CUDA_AVAILABLE else torch.float32
+    
+    # Setup generator
+    generator = torch.Generator(device=device)
+    if seed is None:
+        seed = random.randint(0, 2**31 - 1)
+    generator.manual_seed(seed)
+    
+    logger.info(f"Generating {width}x{height} background: '{prompt}' (seed: {seed})")
+    
+    try:
+        # Choose pipeline based on whether we have an init image
+        if init_image_path and os.path.exists(init_image_path):
+            # Image-to-image pipeline
+            logger.info("Using img2img pipeline")
+            pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+                "runwayml/stable-diffusion-v1-5",
+                torch_dtype=torch_dtype,
+                safety_checker=None,
+                requires_safety_checker=False
+            ).to(device)
+            
+            # Enable memory efficient attention if available
+            try:
+                pipe.enable_attention_slicing()
+                if hasattr(pipe, 'enable_model_cpu_offload'):
+                    pipe.enable_model_cpu_offload()
+            except Exception:
+                pass
+            
+            # Load and resize init image
+            init_image = Image.open(init_image_path).convert("RGB")
+            init_image = init_image.resize((width, height), Image.LANCZOS)
+            
+            # Generate
+            result = pipe(
+                prompt=prompt,
+                image=init_image,
+                strength=0.6,
+                num_inference_steps=num_steps,
+                guidance_scale=guidance_scale,
+                generator=generator,
+                height=height,
+                width=width
+            ).images[0]
+            
+        else:
+            # Text-to-image pipeline
+            logger.info("Using txt2img pipeline")
+            pipe = StableDiffusionPipeline.from_pretrained(
+                "runwayml/stable-diffusion-v1-5",
+                torch_dtype=torch_dtype,
+                safety_checker=None,
+                requires_safety_checker=False
+            ).to(device)
+            
+            # Enable memory efficient attention if available
+            try:
+                pipe.enable_attention_slicing()
+                if hasattr(pipe, 'enable_model_cpu_offload'):
+                    pipe.enable_model_cpu_offload()
+            except Exception:
+                pass
+            
+            # Generate
+            result = pipe(
+                prompt=prompt,
+                height=height,
+                width=width,
+                num_inference_steps=num_steps,
+                guidance_scale=guidance_scale,
+                generator=generator
+            ).images[0]
+        
+        # Save result
+        output_path = TEMP_DIR / f"ai_bg_{int(time.time())}_{seed:08x}.jpg"
+        result.save(output_path, quality=95, optimize=True)
+        
+        # Cleanup GPU memory
+        try:
+            del pipe
+            if TORCH_AVAILABLE and CUDA_AVAILABLE:
+                torch.cuda.empty_cache()
+        except Exception:
+            pass
+        
+        logger.info(f"AI background generated: {output_path}")
+        return str(output_path)
+        
+    except Exception as e:
+        logger.error(f"AI background generation failed: {e}")
+        raise RuntimeError(f"Background generation failed: {e}")
+
 
 # ==============================================================================
 # MAIN PROCESSING PIPELINE