Upload 11 files

README.md

@@ -1,14 +1,68 @@
----
-title: StyleGAN
-emoji: 📊
-colorFrom: purple
-colorTo: red
-sdk: gradio
-sdk_version: 6.9.0
-app_file: app.py
-pinned: false
-license: apache-2.0
-short_description: this a Style GAN V1 Model For High Quality Face Generation.
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: StyleGAN v1
+emoji: 🎨
+colorFrom: blue
+colorTo: purple
+sdk: gradio
+sdk_version: "4.0.0"
+app_file: app.py
+pinned: false
+license: mit
+---
+
+# StyleGAN v1 – Image generation
+
+Generate images with your trained StyleGAN v1 generator.
+
+## How to use
+
+1. **Number of images**: Choose how many images to generate (1–16).
+2. **Resolution**: Select output resolution (4 up to 256).
+3. **Random seed** (optional): Set a seed for reproducible samples.
+4. Click **Generate**.
+
+## Adding your weights to this Space
+
+Your trained weights must be in this repo so the app can load them.
+
+### Option A: Keras weights (recommended)
+
+1. Save your generator in training with:
+   ```python
+   generator.save_weights("weights/generator.weights.h5")
+   ```
+2. Create a `weights` folder in this repo and upload `generator.weights.h5` there.
+3. Or push the file with Git LFS if it is large:
+   ```bash
+   git lfs install
+   git lfs track "weights/*.h5"
+   git add weights/generator.weights.h5 .gitattributes
+   git commit -m "Add generator weights"
+   git push
+   ```
+
+### Option B: State checkpoint (gen_state)
+
+If your training script saves a checkpoint with `ema_trainable` and `non_trainable` (e.g. with `pickle`):
+
+1. Save it as `weights/gen_state.pkl`.
+2. Add the `weights` folder and `gen_state.pkl` to this repo (or use Git LFS for large files).
+
+## Run locally
+
+```bash
+pip install -r requirements.txt
+# Add your weights to weights/generator.weights.h5 (or weights/gen_state.pkl)
+gradio app.py
+```
+
+Open http://localhost:7860 in your browser.
+
+## Project structure
+
+- `app.py` – Gradio UI and entrypoint for the Space
+- `inference.py` – Loads weights and runs generation
+- `loading_model.py` – Builds the StyleGAN generator
+- `layers.py` – Generator/discriminator layers
+- `configuration.py` – Resolutions and hyperparameters
+- `weights/` – Put your generator weights here

app.py

@@ -0,0 +1,105 @@
+"""
+Gradio app for StyleGAN v1 image generation.
+Run: gradio app.py
+For Hugging Face Spaces, the Space runs this file.
+"""
+import os
+import gradio as gr
+
+# Backend must be set before any Keras/JAX imports
+os.environ["KERAS_BACKEND"] = "jax"
+
+from inference import load_weights, generate_images, WEIGHTS_PATH, STATE_PATH
+from configuration import RESOLUTIONS
+
+# Load model once at startup
+_model_loaded = False
+
+
+def load_model_once(resolution=256):
+    global _model_loaded
+    if _model_loaded:
+        return True
+    ok = load_weights(resolution=resolution)
+    if ok:
+        _model_loaded = True
+    return ok
+
+
+def generate(n_images, resolution, seed):
+    n_images = max(1, min(int(n_images), 16))
+    resolution = int(resolution)
+    seed = int(seed) if seed is not None and str(seed).strip() else None
+
+    if not load_model_once(resolution):
+        return (
+            [],
+            "No weights found. Add generator weights to the repo:\n"
+            f"• Keras: {WEIGHTS_PATH}\n"
+            f"• Or state dict: {STATE_PATH}\n"
+            "See README for Hugging Face setup.",
+        )
+
+    try:
+        images = generate_images(n_images, resolution=resolution, seed=seed)
+        if not images:
+            return [], "No images generated."
+        return images, f"Generated {len(images)} image(s) at {resolution}×{resolution}."
+    except Exception as e:
+        import traceback
+        return [], f"Error: {e}\n{traceback.format_exc()}"
+
+
+def main():
+    resolutions = [int(r) for r in RESOLUTIONS]
+
+    with gr.Blocks(
+        title="StyleGAN v1",
+        theme=gr.themes.Soft(),
+    ) as demo:
+        gr.Markdown("# StyleGAN v1 – Image generation")
+        gr.Markdown("Choose how many images to generate and the resolution. Then click **Generate**.")
+
+        with gr.Row():
+            n_images = gr.Slider(
+                minimum=1,
+                maximum=16,
+                value=4,
+                step=1,
+                label="Number of images",
+            )
+            resolution = gr.Dropdown(
+                choices=[str(r) for r in resolutions],
+                value="256",
+                label="Resolution",
+            )
+            seed = gr.Number(
+                value=None,
+                label="Random seed (optional)",
+                precision=0,
+            )
+        with gr.Row():
+            gen_btn = gr.Button("Generate", variant="primary")
+        status = gr.Textbox(label="Status", interactive=False)
+        gallery = gr.Gallery(
+            label="Generated images",
+            show_label=True,
+            columns=4,
+            object_fit="contain",
+            height="auto",
+        )
+
+        gen_btn.click(
+            fn=generate,
+            inputs=[n_images, resolution, seed],
+            outputs=[gallery, status],
+        )
+
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=int(os.environ.get("PORT", 7860)),
+    )
+
+
+if __name__ == "__main__":
+    main()

checkpoint.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d77ad4ceca48e1c0580682e3f336303bda95c33d079b75f887c3dca0818ebd78
+size 190231487

configuration.py

@@ -0,0 +1,23 @@
+RESOLUTIONS = [4, 8, 16, 32, 64, 128, 256]
+
+RES_TO_FILTERS = {
+    4:   512,
+    8:   512,
+    16:  512,
+    32:  512,
+    64:  256,
+    128: 128,
+    256: 64,
+}
+
+LATENT_DIM = 512
+BATCH_SIZE = 32
+LEARNING_RATE = 1e-4
+
+def normalize(img):
+    img = img / 127.5 - 1.0  # Normalize images to [-1,1]
+    return img
+
+def denormalize(img):
+        return (img + 1) / 2
+    

generate.py

@@ -0,0 +1,26 @@
+from imports import jax,jnp,np
+from loading_model import generator
+def inference_and_plot(gen_state, resolution, phase='stable'):
+    rng = gen_state['rng']
+
+    # ── split into 3 keys: next state, z, noise ──────────
+    rng, z_key, noise_key = jax.random.split(rng, 3)
+
+    fixed_z = jax.random.normal(z_key, [16, 512])
+
+    fake_images, _ = generator.stateless_call(
+        gen_state['ema_trainable'],
+        gen_state['non_trainable'],
+        jnp.array(fixed_z),
+        jnp.array(1.0),
+        noise_key,                  # ← dedicated noise key
+    )
+
+    print(f"Fake image range: {float(jnp.min(fake_images)):.3f} to {float(jnp.max(fake_images)):.3f}")
+
+    gen_state = {**gen_state, 'rng': rng}  # ← advance state with consumed rng
+
+    fake_images = (fake_images + 1.0) / 2.0
+    fake_images = jnp.clip(fake_images, 0.0, 1.0)
+    fake_images = np.array(fake_images)
+    return gen_state

imports.py

@@ -0,0 +1,9 @@
+import os
+os.environ["KERAS_BACKEND"] = "jax"  # must be set BEFORE importing keras
+
+import numpy as np
+import keras
+import jax
+import jax.numpy as jnp
+from keras import ops, layers, Model, random, activations
+from keras.models import load_model

inference.py

@@ -0,0 +1,114 @@
+"""
+StyleGAN v1 inference: load weights and generate images.
+Supports Keras .weights.h5 and optional state checkpoint (gen_state).
+"""
+import os
+os.environ["KERAS_BACKEND"] = "jax"
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+
+# Import after backend set
+from loading_model import generator
+from configuration import LATENT_DIM, RESOLUTIONS
+
+# Default paths (override with env or put files in repo)
+WEIGHTS_DIR = os.environ.get("STYLEGAN_WEIGHTS_DIR", "weights")
+WEIGHTS_PATH = os.environ.get("STYLEGAN_WEIGHTS_PATH", os.path.join(WEIGHTS_DIR, "generator.weights.h5"))
+STATE_PATH = os.environ.get("STYLEGAN_STATE_PATH", os.path.join(WEIGHTS_DIR, "gen_state.pkl"))
+
+_gen_state = None
+_use_stateless = False
+
+
+def _build_model():
+    """Build generator with a dummy forward pass so weights can be loaded."""
+    rng = jax.random.PRNGKey(0)
+    dummy_z = jnp.zeros((1, LATENT_DIM))
+    _ = generator(dummy_z, 1.0, rng)
+
+
+def load_weights(resolution=256):
+    """
+    Load generator weights from disk.
+    Prefer WEIGHTS_PATH (Keras .weights.h5). If not found, try STATE_PATH (gen_state pickle).
+    """
+    global _gen_state, _use_stateless
+    _build_model()
+
+    if os.path.isfile(WEIGHTS_PATH):
+        generator.load_weights(WEIGHTS_PATH)
+        _use_stateless = False
+        generator.current_resolution = min(resolution, max(RESOLUTIONS))
+        if resolution not in RESOLUTIONS:
+            generator.current_resolution = max(r for r in RESOLUTIONS if r <= resolution)
+        return True
+
+    try:
+        import pickle
+        if os.path.isfile(STATE_PATH):
+            with open(STATE_PATH, "rb") as f:
+                _gen_state = pickle.load(f)
+            _use_stateless = True
+            return True
+    except Exception:
+        pass
+
+    return False
+
+
+def generate_images(n_images, resolution=256, seed=None):
+    """
+    Generate n_images fake images.
+    Returns list of numpy arrays (H, W, 3) in [0, 255] uint8 for display.
+    """
+    if seed is None:
+        seed = int(np.random.randint(0, 2**31))
+    rng = jax.random.PRNGKey(seed)
+
+    if _use_stateless and _gen_state is not None:
+        return _generate_stateless(n_images, resolution, rng)
+    return _generate_keras(n_images, resolution, rng)
+
+
+def _generate_keras(n_images, resolution, rng):
+    res = min(resolution, max(RESOLUTIONS))
+    res = max(r for r in RESOLUTIONS if r <= res)
+    generator.current_resolution = res
+
+    rng, z_key, noise_key = jax.random.split(rng, 3)
+    z = jax.random.normal(z_key, (n_images, LATENT_DIM))
+
+    out_call = generator(z, 1.0, noise_key)
+    images = out_call[0] if isinstance(out_call, (list, tuple)) else out_call
+    images = (np.array(images) + 1.0) / 2.0
+    images = np.clip(images, 0.0, 1.0)
+    out = []
+    for i in range(n_images):
+        img = (images[i] * 255).astype(np.uint8)
+        out.append(img)
+    return out
+
+
+def _generate_stateless(n_images, resolution, rng):
+    gen_state = _gen_state
+    res = min(resolution, max(RESOLUTIONS))
+    res = max(r for r in RESOLUTIONS if r <= res)
+    # If state was saved with a resolution, we use generator.current_resolution from state
+    # Otherwise we'd need to store it in gen_state; for now use generator attribute
+    generator.current_resolution = res
+
+    rng, z_key, noise_key = jax.random.split(rng, 3)
+    z = jax.random.normal(z_key, (n_images, LATENT_DIM))
+
+    fake_images, _ = generator.stateless_call(
+        gen_state["ema_trainable"],
+        gen_state["non_trainable"],
+        z,
+        1.0,
+        noise_key,
+    )
+    fake_images = (np.array(fake_images) + 1.0) / 2.0
+    fake_images = np.clip(fake_images, 0.0, 1.0)
+    return [(fake_images[i] * 255).astype(np.uint8) for i in range(n_images)]

inspect_checkpoint.py

@@ -0,0 +1,111 @@
+"""
+Inspect a training checkpoint and write a cleaned inference-only state.
+- Drops optimizer variables.
+- If ema_trainable exists, keeps only ema_trainable + non_trainable (drops trainable).
+- Saves to weights/gen_state.pkl for use by inference.py / app.py.
+"""
+import os
+import sys
+import pickle
+
+def _describe(obj, depth=0, max_depth=3):
+    """Return a short description of the object for printing."""
+    if depth > max_depth:
+        return "..."
+    if hasattr(obj, "shape"):
+        return f"array shape={obj.shape} dtype={getattr(obj, 'dtype', '?')}"
+    if isinstance(obj, dict):
+        return "dict(" + ", ".join(f"{k!r}" for k in list(obj.keys())[:8]) + ("..." if len(obj) > 8 else "") + ")"
+    if isinstance(obj, (list, tuple)):
+        return f"{type(obj).__name__} len={len(obj)}"
+    return type(obj).__name__
+
+
+def inspect_checkpoint(path):
+    """Load checkpoint and print top-level keys and a short description of each value."""
+    with open(path, "rb") as f:
+        data = pickle.load(f)
+
+    if not isinstance(data, dict):
+        print(f"Top-level type: {type(data)}")
+        return data
+
+    print("Checkpoint keys and value summary:")
+    print("-" * 60)
+    for k in sorted(data.keys()):
+        v = data[k]
+        desc = _describe(v)
+        print(f"  {k!r}: {desc}")
+    print("-" * 60)
+    return data
+
+
+# Keys we consider optimizer / training-only (to drop)
+OPT_OR_TRAINING_KEYS = {
+    "opt_state", "optimizer", "optimizer_state", "opt",
+    "step", "steps", "epoch", "epochs",
+    "trainable",  # dropped when ema_trainable exists
+    "rng",  # inference creates its own; optional to keep for reproducibility
+}
+
+
+def clean_checkpoint(data, keep_rng=False):
+    """
+    Build a state dict for inference only.
+    - Drop any key in OPT_OR_TRAINING_KEYS (and 'trainable' if ema exists).
+    - If 'ema_trainable' exists, drop 'trainable'.
+    - Keep only ema_trainable (or trainable) + non_trainable.
+    """
+    has_ema = "ema_trainable" in data
+    out = {}
+
+    if has_ema:
+        out["ema_trainable"] = data["ema_trainable"]
+        # drop trainable (we use ema only)
+    else:
+        if "trainable" in data:
+            out["ema_trainable"] = data["trainable"]  # inference expects key "ema_trainable"
+        # else no trainable params in checkpoint
+
+    if "non_trainable" in data:
+        out["non_trainable"] = data["non_trainable"]
+
+    if keep_rng and "rng" in data:
+        out["rng"] = data["rng"]
+
+    return out
+
+
+def main():
+    import argparse
+    p = argparse.ArgumentParser(description="Inspect checkpoint and optionally write cleaned gen_state.pkl")
+    p.add_argument("checkpoint", nargs="?", default="weights/checkpoint.pkl",
+                   help="Path to full training checkpoint (default: weights/checkpoint.pkl)")
+    p.add_argument("-o", "--output", default="weights/gen_state.pkl",
+                   help="Output path for cleaned state (default: weights/gen_state.pkl)")
+    p.add_argument("--no-write", action="store_true", help="Only inspect, do not write cleaned state")
+    p.add_argument("--keep-rng", action="store_true", help="Keep rng in cleaned state")
+    args = p.parse_args()
+
+    if not os.path.isfile(args.checkpoint):
+        print(f"Checkpoint not found: {args.checkpoint}", file=sys.stderr)
+        sys.exit(1)
+
+    data = inspect_checkpoint(args.checkpoint)
+
+    if not isinstance(data, dict):
+        print("Checkpoint is not a dict; cannot clean.", file=sys.stderr)
+        sys.exit(1)
+
+    cleaned = clean_checkpoint(data, keep_rng=args.keep_rng)
+    print("Cleaned state keys:", list(cleaned.keys()))
+
+    if not args.no_write:
+        os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
+        with open(args.output, "wb") as f:
+            pickle.dump(cleaned, f)
+        print(f"Wrote {args.output}")
+
+
+if __name__ == "__main__":
+    main()

layers.py

@@ -0,0 +1,317 @@
+from imports import *          # gets all the libs
+from configuration import *  # gets all the constants
+
+class EqualizedConv2D(keras.layers.Layer):
+    def __init__(self, filters, kernel_size, gain=2.0, **kwargs): 
+        super().__init__(**kwargs)
+        self.filters     = filters
+        self.kernel_size = kernel_size
+        self.gain        = gain
+
+    def build(self, input_shape):
+        self.kernel = self.add_weight(
+            shape=(self.kernel_size, self.kernel_size,
+                   input_shape[-1], self.filters),
+            initializer='random_normal',
+            trainable=True
+        )
+        self.bias = self.add_weight(
+            shape=(self.filters,),
+            initializer='zeros',
+            trainable=True
+        )
+        fan_in = self.kernel_size * self.kernel_size * input_shape[-1]
+        self.c = float(np.sqrt(self.gain / fan_in))
+
+    def call(self, x):
+        x = keras.ops.conv(
+            x, self.kernel * self.c,
+            strides=1, padding='SAME', data_format='channels_last'
+        )
+        return x + self.bias
+
+
+class PixelNorm(keras.layers.Layer):
+    def __init__(self, epsilon=1e-8):
+        super().__init__()
+        self.epsilon = epsilon
+
+    def call(self, x):
+        x_sq    = keras.ops.square(x)
+        mean_sq = keras.ops.mean(x_sq, axis=-1, keepdims=True)
+        x_norm  = x / keras.ops.sqrt(mean_sq + self.epsilon)
+        return x_norm
+
+
+class EqualizedDense(keras.layers.Layer):
+    def __init__(self, units, gain=2.0, lr_multiplier=1.0, **kwargs):
+        super().__init__(**kwargs)
+        self.units          = units
+        self.gain           = gain
+        self.lr_multiplier  = lr_multiplier
+
+    def build(self, input_shape):
+        self.kernel = self.add_weight(
+            shape=(input_shape[-1], self.units),
+            initializer=keras.initializers.RandomNormal(
+                mean=0.0, stddev=1.0 / self.lr_multiplier
+            ),
+            trainable=True
+        )
+        self.bias = self.add_weight(
+            shape=(self.units,),
+            initializer='zeros',
+            trainable=True
+        )
+        fan_in = input_shape[-1]
+        self.c = float(np.sqrt(self.gain / fan_in))
+
+    def call(self, x):
+        return (keras.ops.matmul(x, self.kernel * self.c) + self.bias) * self.lr_multiplier
+
+class MinibatchStd(keras.layers.Layer):
+    def __init__(self, group_size=4, epsilon=1e-8):
+        super().__init__()
+        self.group_size = group_size
+        self.epsilon    = epsilon
+
+    def call(self, x):
+        batch = x.shape[0]                                     # ← static shape not tracer
+        h     = x.shape[1]
+        w     = x.shape[2]
+        c     = x.shape[3]
+
+        group_size = min(self.group_size, batch)               # ← plain Python min
+
+        y    = keras.ops.reshape(x, [group_size, -1, h, w, c])
+        mean = keras.ops.mean(y, axis=0, keepdims=True)
+        var  = keras.ops.mean(keras.ops.square(y - mean), axis=0)
+        std  = keras.ops.sqrt(var + self.epsilon)
+
+        avg_std = keras.ops.mean(std, axis=[1, 2, 3], keepdims=True)
+        avg_std = keras.ops.tile(avg_std, [group_size, h, w, 1])
+
+        return keras.ops.concatenate([x, avg_std], axis=-1)
+        
+# ── gain=1.0 for linear projections ──────────────────────
+def toRGB():
+    return EqualizedConv2D(filters=3, kernel_size=1, gain=1.0)  # no activation
+
+def fromRGB(filters):
+    return keras.Sequential([
+        EqualizedConv2D(filters, kernel_size=1, gain=2.0),
+        keras.layers.LeakyReLU(0.2),                            # activation here
+    ])
+    
+    
+class GenBlock(keras.layers.Layer):
+    def __init__(self, filters):
+        super().__init__()
+        self.conv1 = EqualizedConv2D(filters, kernel_size=4) 
+        self.conv2 = EqualizedConv2D(filters, kernel_size=4)
+        self.act   = keras.layers.LeakyReLU(0.2)
+        self.pnorm = PixelNorm()
+    def call(self, x):
+        x = self.act(self.conv1(x))
+        x = self.pnorm(x)
+        x = self.act(self.conv2(x))
+        x = self.pnorm(x)
+        return x
+
+
+class DiscBlock(keras.layers.Layer):
+    def __init__(self, in_filters, out_filters, **kwargs):
+        super().__init__(**kwargs)
+        self.conv1 = EqualizedConv2D(in_filters, 4)
+        self.conv2 = EqualizedConv2D(out_filters, 4)
+        self.act   = keras.layers.LeakyReLU(0.2)
+        self.pool  = keras.layers.AveragePooling2D(pool_size=2)
+
+    def call(self, x):
+        x = self.act(self.conv1(x))
+        x = self.act(self.conv2(x))
+        x = self.pool(x)
+        return x
+    
+def build_mapping_network():
+    return keras.Sequential([
+    EqualizedDense(512, gain=2.0, lr_multiplier=0.01),  # ← 100x slower
+    keras.layers.LeakyReLU(0.2),
+    EqualizedDense(512, gain=2.0, lr_multiplier=0.01),
+    keras.layers.LeakyReLU(0.2),
+    EqualizedDense(512, gain=2.0, lr_multiplier=0.01),
+    keras.layers.LeakyReLU(0.2),
+    EqualizedDense(512, gain=2.0, lr_multiplier=0.01),
+    keras.layers.LeakyReLU(0.2),
+    EqualizedDense(512, gain=2.0, lr_multiplier=0.01),
+    keras.layers.LeakyReLU(0.2),
+    EqualizedDense(512, gain=2.0, lr_multiplier=0.01),
+    keras.layers.LeakyReLU(0.2),
+    EqualizedDense(512, gain=2.0, lr_multiplier=0.01),
+    keras.layers.LeakyReLU(0.2),
+    EqualizedDense(512, gain=2.0, lr_multiplier=0.01),
+    keras.layers.LeakyReLU(0.2),
+], name="mapping_network")
+    
+class AdaIN(keras.layers.Layer):
+    def __init__(self, channels, w_dim=512):
+        super().__init__()
+        self.channels = channels
+        self.w_dim    = w_dim
+
+        # single affine layer → 2 * channels, then split into ys and yb
+        self.style_transform = EqualizedDense(2 * channels, gain=1.0)
+
+    def call(self, x, w):
+        # ── explicit mean and std over H, W per channel ──────
+        mean = keras.ops.mean(x, axis=[1, 2], keepdims=True)   # (batch, 1, 1, channels)
+        std  = keras.ops.std(x,  axis=[1, 2], keepdims=True) + 1e-8
+
+        normalized_x = (x - mean) / std
+
+        # ── single affine → split into ys and yb ─────────────
+        style = self.style_transform(w)                        # (batch, 2 * channels)
+        ys, yb = keras.ops.split(style, 2, axis=-1)            # each (batch, channels)
+
+        ys = keras.ops.reshape(ys, [-1, 1, 1, self.channels])
+        yb = keras.ops.reshape(yb, [-1, 1, 1, self.channels])
+
+        return ys * normalized_x + yb
+
+# ─── Noise Injection ──────────────────────────────────────
+class NoiseInjection(keras.layers.Layer):
+    def __init__(self, channels):
+        super().__init__()
+        self.B = self.add_weight(
+            shape=(1, 1, 1, channels),
+            initializer="zeros",
+            trainable=True,
+            name="noise_scale"
+        )
+
+    def call(self, x, rng_key):
+        batch = keras.ops.shape(x)[0]
+        h     = keras.ops.shape(x)[1]
+        w     = keras.ops.shape(x)[2]
+        noise = jax.random.normal(rng_key, shape=(batch, h, w, 1))  # ← pure JAX rng
+        return x + self.B * noise
+
+
+# ─── StyleGAN Block ───────────────────────────────────────
+class StyleBlock(keras.layers.Layer):
+    def __init__(self, channels, w_dim=512):
+        super().__init__()
+        self.conv    = EqualizedConv2D(channels, 4)
+        self.noise   = NoiseInjection(channels)
+        self.adain   = AdaIN(channels, w_dim)
+        self.act   = keras.layers.LeakyReLU(0.2)
+
+    def call(self, x, w,rng_key):
+        x = self.conv(x)       # Conv 3x3
+        x = self.noise(x,rng_key)
+        x = self.act(x)
+        x = self.adain(x, w)   # AdaIN with style from w
+        return x
+    
+class StyleGAN_Generator(keras.Model):
+    def __init__(self):
+        super().__init__()
+
+        self.mapping_network = build_mapping_network()
+
+        self.const = self.add_weight(
+            shape=(1, 4, 4, 512),
+            initializer="ones",
+            trainable=True,
+            name="const"
+        )
+
+        # ── register via setattr so Keras tracks them ─────
+        for res in RESOLUTIONS:
+            setattr(self, f"block_{res}_0", StyleBlock(RES_TO_FILTERS[res]))
+            setattr(self, f"block_{res}_1", StyleBlock(RES_TO_FILTERS[res]))
+            setattr(self, f"to_rgb_{res}",  toRGB())
+
+        self.upsample           = keras.layers.UpSampling2D(size=2,interpolation="bilinear")
+        self.current_resolution = 4
+
+    def call(self, z, alpha, rng_key):
+        
+        w = self.mapping_network(z)
+
+        batch = keras.ops.shape(z)[0]
+        x     = keras.ops.tile(self.const, [batch, 1, 1, 1])
+
+        rng_key, subkey = jax.random.split(rng_key)
+        x = getattr(self, "block_4_0")(x, w, subkey)
+        rng_key, subkey = jax.random.split(rng_key)
+        x = getattr(self, "block_4_1")(x, w, subkey)
+
+        if self.current_resolution == 4:
+            return getattr(self, "to_rgb_4")(x)
+
+        for res in RESOLUTIONS[1:]:
+            x_prev = x
+            x      = self.upsample(x)
+            rng_key, subkey = jax.random.split(rng_key)
+            x = getattr(self, f"block_{res}_0")(x, w, subkey)
+            rng_key, subkey = jax.random.split(rng_key)
+            x = getattr(self, f"block_{res}_1")(x, w, subkey)
+
+            if res == self.current_resolution:
+                old_rgb = getattr(self, f"to_rgb_{res // 2}")(self.upsample(x_prev))
+                new_rgb = getattr(self, f"to_rgb_{res}")(x)
+                return (1 - alpha) * old_rgb + alpha * new_rgb
+
+        return getattr(self, f"to_rgb_{self.current_resolution}")(x)
+
+
+class ProGAN_Discriminator(keras.Model):
+    def __init__(self):
+        super().__init__()
+
+        # ── register via setattr so Keras tracks them ─────
+        for res in RESOLUTIONS[1:]:
+            setattr(self, f"block_{res}", DiscBlock(RES_TO_FILTERS[res], RES_TO_FILTERS[res // 2]))
+        for res in RESOLUTIONS:
+            setattr(self, f"from_rgb_{res}", fromRGB(RES_TO_FILTERS[res]))
+
+        self.minibatch_std = MinibatchStd()
+        self.final_conv    = EqualizedConv2D(512, kernel_size=3)
+        self.final_dense_1 = EqualizedDense(512)        # ← missing intermediate dense
+        self.final_dense_2 = EqualizedDense(1)
+        self.flatten       = keras.layers.Flatten()
+        self.act           = keras.layers.LeakyReLU(0.2)
+        self.downsample    = keras.layers.AveragePooling2D(pool_size=2)
+        self.current_resolution = 4
+
+    def call(self, img, alpha):
+        if self.current_resolution == 4:
+            x = getattr(self, "from_rgb_4")(img)
+            x = self.minibatch_std(x)
+            x = self.act(self.final_conv(x))
+            x = self.flatten(x)
+            x = self.act(self.final_dense_1(x))        # ← intermediate dense
+            return self.final_dense_2(x)
+
+        cur_res  = self.current_resolution
+        prev_res = cur_res // 2
+
+        x_new = getattr(self, f"from_rgb_{cur_res}")(img)
+        x_new = getattr(self, f"block_{cur_res}")(x_new)
+
+        x_old = self.downsample(img)
+        x_old = getattr(self, f"from_rgb_{prev_res}")(x_old)
+
+        x = (1 - alpha) * x_old + alpha * x_new
+
+        for res in reversed(RESOLUTIONS[1:]):
+            if res >= cur_res:
+                continue
+            x = getattr(self, f"block_{res}")(x)
+
+        x = self.minibatch_std(x)
+        x = self.act(self.final_conv(x))
+        x = self.flatten(x)
+        x = self.act(self.final_dense_1(x))            # ← intermediate dense
+        return self.final_dense_2(x)

loading_model.py

@@ -0,0 +1,4 @@
+from imports import *          # gets all the libs
+from layers import *
+# ─── build models ────────────────────────────────────────
+generator     = StyleGAN_Generator()

requirements.txt

@@ -0,0 +1,6 @@
+# StyleGAN v1 – Gradio + Hugging Face Space
+gradio>=4.0.0
+keras>=3.0.0
+jax>=0.4.0
+jaxlib>=0.4.0
+numpy>=1.24.0,<2.0.0