[Update]: Refined app.py for enhanced GPU handling and memory operations 🌟

- Added: Functions to initialize device settings and manage GPU optimizations for Stable Diffusion.
- Updated: MemoryWave and EmotionalContext classes to support dynamic device management.
- Improved: Gradio interface to utilize GPU capabilities more effectively during memory operations.
- Pro Tip of the Commit: When it comes to memory waves, let’s keep the flow smooth and the devices ready! 🌊🚀
Aye, Aye! 🚢

app.py

@@ -10,7 +10,6 @@ and interact like waves in an ocean of consciousness.
 
 Created by: Aye & Hue (with Trisha from Accounting keeping the numbers flowing)
 """
-import spaces
 import os
 import gradio as gr
 import torch
@@ -39,58 +38,67 @@ EMOTION_RANGE = (-5, 5)  # Range for emotional valence
 AROUSAL_RANGE = (0, 255)  # Range for arousal
 MAX_SEED = 999999999  # Maximum seed value for art generation
 
-# Configure device - Use Zero GPU if available
-if torch.cuda.is_available():
-    print("🚀 Zero GPU detected! Enabling hardware acceleration...")
-    device = "cuda"
-    # Enable Zero GPU optimizations
-    torch.backends.cuda.matmul.allow_tf32 = True
-    torch.backends.cudnn.allow_tf32 = True
-    torch.backends.cudnn.benchmark = True
-else:
-    print("⚡ Running on CPU")
-    device = "cpu"
-
-# Try to import Stable Diffusion components
+# Initialize everything for GPU by default
+device = "cuda" if torch.cuda.is_available() else "cpu"
 STABLE_DIFFUSION_AVAILABLE = False
 pipe = None
-try:
-    from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
-    STABLE_DIFFUSION_AVAILABLE = True
+
+def get_device():
+    """Get the appropriate device for the current context."""
+    if torch.cuda.is_available():
+        return "cuda"
+    print("⚠️ Warning: CUDA not available, falling back to CPU")
+    return "cpu"
+
+def init_stable_diffusion(device_str: str = None):
+    """Initialize Stable Diffusion optimized for GPU usage."""
+    global STABLE_DIFFUSION_AVAILABLE, pipe
+    
+    # Always try to use CUDA if available
+    if device_str is None:
+        device_str = get_device()
     
-    # Initialize Stable Diffusion pipeline with Zero GPU optimizations
     try:
-        model_id = "stabilityai/stable-diffusion-xl-base-1.0"
-        print(f"Loading Stable Diffusion model: {model_id}")
+        from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
+        STABLE_DIFFUSION_AVAILABLE = True
         
-        # Always use float32 for CPU, only use float16 for GPU
-        model_dtype = torch.float16 if device == "cuda" else torch.float32
-        
-        pipe = DiffusionPipeline.from_pretrained(
-            model_id,
-            torch_dtype=model_dtype,
-            use_safetensors=True,
-            variant="fp16" if device == "cuda" else None
-        )
-        pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config)
-        pipe.to(device)
-        
-        # Enable memory optimizations only on GPU
-        if device == "cuda":
-            pipe.enable_model_cpu_offload()
-            pipe.enable_vae_slicing()
-            pipe.enable_vae_tiling()
-            pipe.enable_attention_slicing(slice_size="max")
-            print("✨ Stable Diffusion loaded with GPU optimizations")
-        else:
-            print("✨ Stable Diffusion loaded in CPU mode")
+        try:
+            model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+            print(f"🚀 Loading Stable Diffusion model: {model_id}")
             
-    except Exception as e:
-        print(f"⚠️ Warning: Failed to initialize Stable Diffusion: {e}")
-        STABLE_DIFFUSION_AVAILABLE = False
-        pipe = None
-except ImportError:
-    print("⚠️ Warning: diffusers package not available. Artistic visualization will be disabled.")
+            if device_str == "cuda":
+                print("🎮 Initializing with GPU optimizations...")
+                pipe = DiffusionPipeline.from_pretrained(
+                    model_id,
+                    torch_dtype=torch.float16,
+                    use_safetensors=True,
+                    variant="fp16"
+                )
+                pipe.to(device_str)
+                # Enable all GPU optimizations
+                pipe.enable_model_cpu_offload()
+                pipe.enable_vae_slicing()
+                pipe.enable_vae_tiling()
+                pipe.enable_attention_slicing(slice_size="max")
+                print("✨ Stable Diffusion loaded with full GPU optimizations")
+            else:
+                print("⚠️ CUDA not available - Stable Diffusion may be slow")
+                pipe = DiffusionPipeline.from_pretrained(
+                    model_id,
+                    torch_dtype=torch.float32,
+                    use_safetensors=True
+                )
+                pipe.to("cpu")
+            
+            # Set scheduler
+            pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config)
+                
+        except Exception as e:
+            print(f"❌ Failed to initialize Stable Diffusion: {e}")
+            STABLE_DIFFUSION_AVAILABLE = False
+            pipe = None
+    except ImportError:
+        print("❌ diffusers package not available. Artistic visualization will be disabled.")
 
 # Create a directory for memory snapshots if it doesn't exist
 MEMORY_DIR = "memory_snapshots"
@@ -106,12 +114,12 @@ class EmotionalContext:
         context (torch.Tensor): Contextual flags (16-bit in paper)
         safety (torch.Tensor): Psychological safety indicator
     """
-    def __init__(self, device="cuda" if torch.cuda.is_available() else "cpu"):
-        self.device = device
-        self.valence = torch.zeros(1).to(device)  # -128 to 127: negative to positive
-        self.arousal = torch.zeros(1).to(device)  # 0 to 255: intensity level
-        self.context = torch.zeros(1).to(device)  # Contextual flags
-        self.safety = torch.ones(1).to(device)    # Psychological safety indicator
+    def __init__(self, device_str="cpu"):
+        self.device = device_str
+        self.valence = torch.zeros(1, device=device_str)  # -128 to 127: negative to positive
+        self.arousal = torch.zeros(1, device=device_str)  # 0 to 255: intensity level
+        self.context = torch.zeros(1, device=device_str)  # Contextual flags
+        self.safety = torch.ones(1, device=device_str)    # Psychological safety indicator
         
         # Track emotional history for visualization
         self.history = {
@@ -162,6 +170,15 @@ class EmotionalContext:
         """String representation of emotional context."""
         return f"EmotionalContext(valence={self.valence.item():.2f}, arousal={self.arousal.item():.2f})"
 
+    def to(self, device_str):
+        """Move the context to a different device."""
+        self.device = device_str
+        self.valence = self.valence.to(device_str)
+        self.arousal = self.arousal.to(device_str)
+        self.context = self.context.to(device_str)
+        self.safety = self.safety.to(device_str)
+        return self
+
 class MemoryWave:
     """
     Implements the wave-based memory patterns from Mem|8 paper.
@@ -171,26 +188,26 @@ class MemoryWave:
     """
     def __init__(self, 
                  size: int = DEFAULT_GRID_SIZE, 
-                 device: str = device):  # Use global device
+                 device_str: str = "cpu"):  # Use global device
         """
         Initialize a memory wave system.
         
         Args:
             size: Size of the memory grid (NxN)
-            device: Device to use for computations
+            device_str: Device to use for computations
         """
         self.size = size
-        self.device = device
-        self.grid = torch.zeros((size, size), device=device)
-        self.emotion = EmotionalContext(device)
+        self.device = device_str
+        self.grid = torch.zeros((size, size), device=device_str)
+        self.emotion = EmotionalContext(device_str)
         
         # Initialize coordinates for wave calculations
-        self.x = torch.linspace(0, 2*np.pi, size, device=device)
-        self.y = torch.linspace(0, 2*np.pi, size, device=device)
+        self.x = torch.linspace(0, 2*np.pi, size, device=device_str)
+        self.y = torch.linspace(0, 2*np.pi, size, device=device_str)
         self.X, self.Y = torch.meshgrid(self.x, self.y, indexing='ij')
         
         # Memory storage for different types
-        self.memory_types = {i: torch.zeros((size, size), device=device) for i in range(6)}
+        self.memory_types = {i: torch.zeros((size, size), device=device_str) for i in range(6)}
         
         # History of wave states for animation
         self.history = []
@@ -488,7 +505,7 @@ class MemoryWave:
                 "arousal": self.emotion.arousal.item(),
             }
         }
-    @spaces.GPU(duration=10)
+
     def generate_memory_reconstruction(self,
                                       emotion_valence: float,
                                       corruption_level: float = 0.3) -> Dict:
@@ -735,51 +752,8 @@ class MemoryWave:
             
         return filepath
 
-def generate_memory_prompt(operation: str, emotion_valence: float) -> str:
-    """Generate an artistic prompt based on the memory operation and emotional context."""
-    
-    # Base prompts for each operation type
-    operation_prompts = {
-        "wave_memory": "A serene ocean of consciousness with rippling waves of memory, ",
-        "interference": "Multiple waves of thought intersecting and creating intricate patterns, ",
-        "resonance": "Harmonious waves of memory resonating with emotional energy, ",
-        "reconstruction": "Fragments of memory waves reforming into a complete pattern, ",
-        "hot_tub": "A safe sanctuary of gentle memory waves with healing energy, "
-    }
-    
-    # Emotional modifiers based on valence
-    if emotion_valence < -3:
-        emotion_desc = "dark and turbulent, with deep indigo and violet hues, expressing profound melancholy"
-    elif emotion_valence < -1:
-        emotion_desc = "muted and somber, with cool blues and grays, showing gentle sadness"
-    elif emotion_valence < 1:
-        emotion_desc = "balanced and neutral, with soft pastels, reflecting calm contemplation"
-    elif emotion_valence < 3:
-        emotion_desc = "warm and uplifting, with golden yellows and soft oranges, radiating joy"
-    else:
-        emotion_desc = "brilliant and ecstatic, with vibrant rainbow colors, bursting with happiness"
-    
-    # Artistic style modifiers
-    style = (
-        "digital art in the style of a quantum visualization, "
-        "highly detailed, smooth gradients, "
-        "abstract yet meaningful, "
-        "inspired by neural networks and consciousness"
-    )
-    
-    # Combine all elements
-    base_prompt = operation_prompts.get(operation, operation_prompts["wave_memory"])
-    prompt = f"{base_prompt}{emotion_desc}, {style}"
-    
-    return prompt
-
-@spaces.GPU(duration=30)
-def create_interface():
-    """Create the Gradio interface for the Mem|8 Wave Memory Explorer."""
-    # Initialize with Zero GPU device
-    memory_wave = MemoryWave(device=device)
-    
     def process_memory_operation(
+        self,
         operation: str,
         emotion_valence: float,
         grid_size: int = DEFAULT_GRID_SIZE,
@@ -791,40 +765,40 @@ def create_interface():
         """Process a memory operation and return visualizations."""
         
         # Resize grid if needed
-        if grid_size != memory_wave.size:
-            memory_wave.__init__(size=grid_size)
+        if grid_size != self.size:
+            self.__init__(size=grid_size, device=self.device)
         
         # Process based on operation type
         if operation == "wave_memory":
-            result = memory_wave.generate_wave_memory(emotion_valence)
+            result = self.generate_wave_memory(emotion_valence)
             wave_title = "Wave Memory Pattern"
             wave_data = result["wave"]
             
         elif operation == "interference":
-            result = memory_wave.generate_interference_pattern(emotion_valence)
+            result = self.generate_interference_pattern(emotion_valence)
             wave_title = "Interference Pattern"
             wave_data = result["wave"]
             
         elif operation == "resonance":
-            result = memory_wave.generate_resonance_pattern(emotion_valence)
+            result = self.generate_resonance_pattern(emotion_valence)
             wave_title = "Resonance Pattern"
             wave_data = result["wave"]
             
         elif operation == "reconstruction":
-            result = memory_wave.generate_memory_reconstruction(emotion_valence)
+            result = self.generate_memory_reconstruction(emotion_valence)
             wave_title = "Memory Reconstruction"
             wave_data = result["reconstructed"]
             
         elif operation == "hot_tub":
-            result = memory_wave.generate_hot_tub_simulation(
+            result = self.generate_hot_tub_simulation(
                 emotion_valence, comfort_level, exploration_depth
             )
             wave_title = "Hot Tub Exploration"
             wave_data = result["safe_result"]
         
         # Create visualizations
-        wave_plot = memory_wave.visualize_wave_pattern(wave_data, wave_title)
-        emotion_plot = memory_wave.visualize_emotional_history()
+        wave_plot = self.visualize_wave_pattern(wave_data, wave_title)
+        emotion_plot = self.visualize_emotional_history()
         
         # Generate artistic visualization if requested
         art_output = None
@@ -852,11 +826,92 @@ def create_interface():
         metrics_str += f"• Arousal: {result['emotion']['arousal']:.2f}\n"
         
         # Save memory snapshot
-        snapshot_path = memory_wave.save_memory_snapshot(operation)
+        snapshot_path = self.save_memory_snapshot(operation)
         metrics_str += f"\n💾 Memory snapshot saved: {snapshot_path}"
         
         return metrics_str, wave_plot, emotion_plot, art_output
 
+    def to(self, device_str):
+        """Move the wave system to a different device."""
+        self.device = device_str
+        self.grid = self.grid.to(device_str)
+        self.emotion = self.emotion.to(device_str)
+        self.x = self.x.to(device_str)
+        self.y = self.y.to(device_str)
+        self.X = self.X.to(device_str)
+        self.Y = self.Y.to(device_str)
+        self.memory_types = {k: v.to(device_str) for k, v in self.memory_types.items()}
+        return self
+
+def generate_memory_prompt(operation: str, emotion_valence: float) -> str:
+    """Generate an artistic prompt based on the memory operation and emotional context."""
+    
+    # Base prompts for each operation type
+    operation_prompts = {
+        "wave_memory": "A serene ocean of consciousness with rippling waves of memory, ",
+        "interference": "Multiple waves of thought intersecting and creating intricate patterns, ",
+        "resonance": "Harmonious waves of memory resonating with emotional energy, ",
+        "reconstruction": "Fragments of memory waves reforming into a complete pattern, ",
+        "hot_tub": "A safe sanctuary of gentle memory waves with healing energy, "
+    }
+    
+    # Emotional modifiers based on valence
+    if emotion_valence < -3:
+        emotion_desc = "dark and turbulent, with deep indigo and violet hues, expressing profound melancholy"
+    elif emotion_valence < -1:
+        emotion_desc = "muted and somber, with cool blues and grays, showing gentle sadness"
+    elif emotion_valence < 1:
+        emotion_desc = "balanced and neutral, with soft pastels, reflecting calm contemplation"
+    elif emotion_valence < 3:
+        emotion_desc = "warm and uplifting, with golden yellows and soft oranges, radiating joy"
+    else:
+        emotion_desc = "brilliant and ecstatic, with vibrant rainbow colors, bursting with happiness"
+    
+    # Artistic style modifiers
+    style = (
+        "digital art in the style of a quantum visualization, "
+        "highly detailed, smooth gradients, "
+        "abstract yet meaningful, "
+        "inspired by neural networks and consciousness"
+    )
+    
+    # Combine all elements
+    base_prompt = operation_prompts.get(operation, operation_prompts["wave_memory"])
+    prompt = f"{base_prompt}{emotion_desc}, {style}"
+    
+    return prompt
+
+def create_interface():
+    """Create the Gradio interface for the Mem|8 Wave Memory Explorer."""
+    # Initialize with GPU by default
+    device_str = get_device()
+    memory_wave = MemoryWave(device_str=device_str)
+    
+    def process_with_gpu(*args, **kwargs):
+        """Process operations on GPU with optimized memory handling."""
+        device_str = get_device()
+        
+        try:
+            # Ensure we're on GPU if available
+            if device_str == "cuda" and memory_wave.device != "cuda":
+                memory_wave.to("cuda")
+                if pipe is not None and pipe.device.type != "cuda":
+                    pipe.to("cuda")
+            
+            result = memory_wave.process_memory_operation(*args, **kwargs)
+            return result
+            
+        except torch.cuda.OutOfMemoryError:
+            print("⚠️ GPU out of memory - falling back to CPU")
+            memory_wave.to("cpu")
+            if pipe is not None:
+                pipe.to("cpu")
+            return memory_wave.process_memory_operation(*args, **kwargs)
+            
+        except Exception as e:
+            print(f"❌ Error during processing: {e}")
+            return None
+
     # Create the interface
     with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="blue")) as demo:
         gr.Markdown("""
@@ -939,7 +994,7 @@ def create_interface():
         
         # Set up event handlers
         run_btn.click(
-            process_memory_operation,
+            process_with_gpu,  # Use the wrapped function
             inputs=[
                 operation_input,
                 emotion_input,
@@ -974,6 +1029,9 @@ if __name__ == "__main__":
     # Configure Gradio for Hugging Face Spaces
     demo = create_interface()
     
+    # Initialize Stable Diffusion on CPU first
+    init_stable_diffusion("cpu")
+    
     # Enable queuing for better resource management
     demo.queue(max_size=10)