Update app.py

app.py

@@ -8,6 +8,11 @@ import io
 import time
 
 # ===== CONFIG =====
+# Print debug info
+print(f"PyTorch version: {torch.__version__}")
+print(f"CUDA available: {torch.cuda.is_available()}")
+print(f"CUDA device count: {torch.cuda.device_count()}")
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if device == "cuda" else torch.float32
 
@@ -20,9 +25,21 @@ pipe = DiffusionPipeline.from_pretrained(
 )
 pipe.to(device)
 
-# Enable memory efficient attention and channels last for better performance
-pipe.enable_xformers_memory_efficient_attention()
-pipe.unet.to(memory_format=torch.channels_last)
+# Enable optimizations only if GPU is available
+if device == "cuda":
+    try:
+        pipe.enable_xformers_memory_efficient_attention()
+        print("Enabled xformers memory efficient attention")
+    except Exception as e:
+        print(f"Could not enable xformers: {str(e)}")
+    
+    try:
+        pipe.unet.to(memory_format=torch.channels_last)
+        print("Enabled channels last memory format")
+    except Exception as e:
+        print(f"Could not enable channels last: {str(e)}")
+else:
+    print("Running on CPU - skipping GPU optimizations")
 
 MAX_SEED = np.iinfo(np.int32).max
 IMAGE_SIZE = 1024  # Same as original code
@@ -74,27 +91,34 @@ def generate(
     
     generator = torch.manual_seed(seed)
     
-    # Ultra-fast generation with minimal steps
-    result = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        width=IMAGE_SIZE,
-        height=IMAGE_SIZE,
-        guidance_scale=guidance_scale,
-        num_inference_steps=max(1, num_inference_steps),  # Minimum 1 step
-        generator=generator,
-    ).images[0]
-    
-    # Optimized watermark and JPG conversion
-    watermarked = add_watermark(result)
-    buffer = io.BytesIO()
-    watermarked.save(buffer, format="JPEG", quality=85, optimize=True)
-    buffer.seek(0)
-    
-    gen_time = time.time() - start_time
-    status = f"✔️ Generated in {gen_time:.2f}s | Seed: {seed}"
+    try:
+        # Ultra-fast generation with minimal steps
+        result = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            width=IMAGE_SIZE,
+            height=IMAGE_SIZE,
+            guidance_scale=guidance_scale,
+            num_inference_steps=max(1, num_inference_steps),  # Minimum 1 step
+            generator=generator,
+        ).images[0]
+        
+        # Optimized watermark and JPG conversion
+        watermarked = add_watermark(result)
+        buffer = io.BytesIO()
+        watermarked.save(buffer, format="JPEG", quality=85, optimize=True)
+        buffer.seek(0)
+        
+        gen_time = time.time() - start_time
+        status = f"✔️ Generated in {gen_time:.2f}s | Seed: {seed}"
+        
+        return Image.open(buffer), status
     
-    return Image.open(buffer), status
+    except torch.cuda.OutOfMemoryError:
+        return None, "⚠️ GPU out of memory - try a simpler prompt"
+    except Exception as e:
+        print(f"Generation error: {str(e)}")
+        return None, f"⚠️ Error: {str(e)[:200]}"
 
 # ===== EXAMPLES =====
 examples = [