Update app.py

app.py

@@ -1,154 +1,202 @@
 import gradio as gr
+import tensorflow as tf
+from tensorflow import keras
+from tensorflow.keras import layers
 import numpy as np
-import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
-import torch
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
-
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
-
-    return image, seed
-
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
-
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-
-            run_button = gr.Button("Run", scale=0, variant="primary")
-
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2,  # Replace with defaults that work for your model
-                )
-
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, seed],
-    )
-
-if __name__ == "__main__":
-    demo.launch()
+import matplotlib.pyplot as plt
+import os
+from PIL import Image
+import math
+
+# =========================================
+# 1. SETTINGS
+# =========================================
+IMG_SIZE = 48   # Slightly larger than CIFAR for better detail
+CHANNELS = 3
+TIMESTEPS = 200
+BATCH_SIZE = 32
+
+# =========================================
+# 2. GENERATE CUSTOM "VEDA" DATASET
+# =========================================
+def generate_veda_patterns(num_images=1000):
+    """
+    Creates a dataset of mathematical/spiritual patterns (Mandalas/Energy Orbs)
+    instead of using boring CIFAR-10 images.
+    """
+    print(f"Generating {num_images} unique Veda patterns...")
+    data = []
+    
+    for _ in range(num_images):
+        # Create a grid
+        x = np.linspace(-1, 1, IMG_SIZE)
+        y = np.linspace(-1, 1, IMG_SIZE)
+        X, Y = np.meshgrid(x, y)
+        
+        # Random parameters for the pattern
+        freq = np.random.uniform(2, 10)
+        phase = np.random.uniform(0, np.pi)
+        center_x = np.random.uniform(-0.5, 0.5)
+        center_y = np.random.uniform(-0.5, 0.5)
+        
+        # Math: Radial Gradient + Sine waves (Ripple effect)
+        R = np.sqrt((X - center_x)**2 + (Y - center_y)**2)
+        pattern = np.sin(freq * R - phase)
+        
+        # Colorize (RGB)
+        img = np.zeros((IMG_SIZE, IMG_SIZE, 3))
+        
+        # Random colors
+        r_boost = np.random.rand()
+        g_boost = np.random.rand()
+        b_boost = np.random.rand()
+        
+        img[:, :, 0] = (pattern + 1) / 2 * r_boost # Red
+        img[:, :, 1] = (pattern + 1) / 2 * g_boost # Green
+        img[:, :, 2] = (pattern + 1) / 2 * b_boost # Blue
+        
+        data.append(img)
+        
+    return np.array(data).astype("float32")
+
+# =========================================
+# 3. DIFFUSION MATH
+# =========================================
+beta_np = np.linspace(0.0001, 0.02, TIMESTEPS)
+beta = tf.cast(beta_np, dtype=tf.float32)
+alpha = 1.0 - beta
+alpha_bar = tf.math.cumprod(alpha, axis=0)
+
+def forward_noise(x_0, t):
+    n = tf.shape(x_0)[0]
+    a_bar = tf.gather(alpha_bar, t)
+    a_bar = tf.reshape(a_bar, [n, 1, 1, 1])
+    noise = tf.random.normal(shape=tf.shape(x_0), dtype=tf.float32)
+    return (tf.sqrt(a_bar) * x_0) + (tf.sqrt(1.0 - a_bar) * noise), noise
+
+# =========================================
+# 4. THE BRAIN (U-NET)
+# =========================================
+def build_unet():
+    image_input = layers.Input(shape=(IMG_SIZE, IMG_SIZE, CHANNELS))
+    time_input = layers.Input(shape=(1,))
+    
+    # Embed Time
+    t = layers.Dense(IMG_SIZE, activation="relu")(time_input)
+    t = layers.Reshape((1, 1, IMG_SIZE))(t)
+
+    # In
+    x = layers.Conv2D(32, 3, padding="same", activation="relu")(image_input)
+    
+    # Simple addition of time info
+    # Note: We keep architecture simple for CPU speed
+    x = layers.Conv2D(64, 3, strides=2, padding="same", activation="relu")(x) # 24x24
+    x = layers.Conv2D(128, 3, strides=2, padding="same", activation="relu")(x) # 12x12
+    x = layers.Conv2D(256, 3, padding="same", activation="relu")(x)
+    
+    # Decode
+    x = layers.Conv2DTranspose(128, 3, strides=2, padding="same", activation="relu")(x)
+    x = layers.Conv2DTranspose(64, 3, strides=2, padding="same", activation="relu")(x)
+    x = layers.Conv2D(32, 3, padding="same", activation="relu")(x)
+    
+    outputs = layers.Conv2D(CHANNELS, 1, padding="same")(x)
+    return keras.Model([image_input, time_input], outputs)
+
+model = build_unet()
+optimizer = keras.optimizers.Adam(learning_rate=1e-4)
+loss_fn = keras.losses.MeanSquaredError()
+
+# Load weights if trained
+if os.path.exists("veda_art_model.keras"):
+    try:
+        model.load_weights("veda_art_model.keras")
+        print("Brain Loaded!")
+    except:
+        pass
+
+# =========================================
+# 5. TRAINING FUNCTION
+# =========================================
+def train_model(epochs):
+    # 1. Generate fresh data
+    x_train = generate_veda_patterns(num_images=500) # 500 images is enough for patterns
+    dataset = tf.data.Dataset.from_tensor_slices(x_train).batch(BATCH_SIZE).shuffle(500)
+    
+    yield "Dataset Generated (Mathematical Patterns). Starting Training..."
+    
+    for epoch in range(int(epochs)):
+        total_loss = 0
+        steps = 0
+        for batch in dataset:
+            batch_size = tf.shape(batch)[0]
+            t = tf.random.uniform([batch_size], minval=0, maxval=TIMESTEPS, dtype=tf.int32)
+            noisy_images, noise = forward_noise(batch, t)
+            
+            with tf.GradientTape() as tape:
+                pred_noise = model([noisy_images, t], training=True)
+                loss = loss_fn(noise, pred_noise)
+            
+            grads = tape.gradient(loss, model.trainable_variables)
+            optimizer.apply_gradients(zip(grads, model.trainable_variables))
+            total_loss += loss
+            steps += 1
+            
+        yield f"Epoch {epoch+1}/{epochs} - Loss: {total_loss/steps:.4f}"
+        model.save_weights("veda_art_model.keras")
+        
+    yield "Training Done! Go to 'Dream' tab."
+
+# =========================================
+# 6. GENERATION FUNCTION
+# =========================================
+def generate_art():
+    # Start with static
+    img = tf.random.normal((1, IMG_SIZE, IMG_SIZE, CHANNELS), dtype=tf.float32)
+    
+    # Reverse Diffusion
+    for i in range(TIMESTEPS - 1, 0, -1):
+        t = tf.fill([1], i)
+        pred_noise = model([img, t], training=False)
+        
+        alpha_t = tf.gather(alpha, i)
+        alpha_bar_t = tf.gather(alpha_bar, i)
+        beta_t = tf.gather(beta, i)
+        
+        term1 = 1.0 / tf.sqrt(alpha_t)
+        term2 = (1.0 - alpha_t) / tf.sqrt(1.0 - alpha_bar_t)
+        
+        img = term1 * (img - term2 * pred_noise)
+        
+        if i > 1:
+            z = tf.random.normal((1, IMG_SIZE, IMG_SIZE, CHANNELS), dtype=tf.float32)
+            img = img + (tf.sqrt(beta_t) * z)
+
+    # Normalize for display
+    img = tf.clip_by_value(img, 0.0, 1.0)
+    img = img[0].numpy()
+    img = (img * 255).astype(np.uint8)
+    
+    # Resize to be big and pretty
+    pil_img = Image.fromarray(img)
+    pil_img = pil_img.resize((300, 300), Image.BILINEAR)
+    return pil_img
+
+# =========================================
+# 7. UI
+# =========================================
+with gr.Blocks(title="Veda Art Engine") as demo:
+    gr.Markdown("# Veda Art Engine")
+    gr.Markdown("This AI generates its own training data using math, then learns to dream up new patterns.")
+    
+    with gr.Tab("Dream"):
+        gen_btn = gr.Button("Generate Sacred Pattern", variant="primary")
+        out_img = gr.Image(label="Veda Dream")
+        gen_btn.click(generate_art, outputs=out_img)
+        
+    with gr.Tab("Train"):
+        gr.Markdown("Click to generate a dataset of mathematical mandalas and train the AI on them.")
+        train_btn = gr.Button("Train AI (5 Epochs)")
+        log = gr.Textbox(label="Status")
+        train_btn.click(lambda: train_model(5), outputs=log)
+
+demo.launch()