Upload 9 files

app.py

@@ -1,105 +1,55 @@
-"""
-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
+from inference import load_weights, generate_images, WEIGHTS_H5, CHECKPOINT_PKL
 
-# Load model once at startup
 _model_loaded = False
 
-
-def load_model_once(resolution=256):
+def load_model_once():
     global _model_loaded
     if _model_loaded:
         return True
-    ok = load_weights(resolution=resolution)
+    ok = load_weights(resolution=256)
     if ok:
         _model_loaded = True
     return ok
 
-
-def generate(n_images, resolution, seed):
+def generate(n_images, 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.",
-        )
+    if not load_model_once():
+        return [], f"No weights found. Put checkpoint.pkl in:\n{os.path.dirname(WEIGHTS_H5)}"
 
     try:
-        images = generate_images(n_images, resolution=resolution, seed=seed)
-        if not images:
+        images = generate_images(n_images, resolution=256, seed=seed)
+        if images is None or len(images) == 0:
             return [], "No images generated."
-        return images, f"Generated {len(images)} image(s) at {resolution}×{resolution}."
+        return images, f"Generated {len(images)} image(s) at 256×256."
     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 – Face Generation")
+        gr.Markdown("Generate faces at 256×256. Optionally set a seed for reproducibility.")
 
-    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",
-        )
+            n_images = gr.Slider(minimum=1, maximum=16, value=4, step=1, label="Number of images")
+            seed     = gr.Number(value=None, label="Random seed (optional)", precision=0)
 
-        gen_btn.click(
-            fn=generate,
-            inputs=[n_images, resolution, seed],
-            outputs=[gallery, status],
-        )
+        gen_btn = gr.Button("Generate", variant="primary")
+        status  = gr.Textbox(label="Status", interactive=False)
+        gallery = gr.Gallery(label="Generated images", columns=4, object_fit="contain", height="auto")
+
+        gen_btn.click(fn=generate, inputs=[n_images, seed], outputs=[gallery, status])
 
     demo.launch(
         server_name="0.0.0.0",
         server_port=int(os.environ.get("PORT", 7860)),
     )
 
-
 if __name__ == "__main__":
-    main()
+    main()

check.py

@@ -0,0 +1,17 @@
+import joblib
+import jax.numpy as jnp
+from loading_model import generator
+from inference import _build_model
+
+_build_model()
+
+# Load checkpoint
+ck = joblib.load("weights/checkpoint.pkl")
+
+print("=== CHECKPOINT variables ===")
+for i, w in enumerate(ck["gen_trainable"]):
+    print(f"  [{i:03d}] shape={str(w.shape):30s}")
+
+print("\n=== MODEL variables ===")
+for i, v in enumerate(generator.trainable_variables):
+    print(f"  [{i:03d}] shape={str(v.shape):30s}  name={v.name}")

inference.py

@@ -1,7 +1,3 @@
-"""
-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"
 
@@ -9,106 +5,122 @@ 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"))
+WEIGHTS_DIR    = os.path.join(os.path.dirname(os.path.abspath(__file__)), "weights")
+WEIGHTS_H5     = os.path.join(WEIGHTS_DIR, "generator.weights.h5")
+CHECKPOINT_PKL = os.path.join(WEIGHTS_DIR, "checkpoint.pkl")
 
-_gen_state = None
+_gen_state     = None
 _use_stateless = False
+_saved_alpha   = 1.0
+
+TARGET_RES     = 256   # hardcoded — only generate at 256
+
+
+def _postprocess(images: np.ndarray) -> list:
+    raw_min, raw_max = images.min(), images.max()
+    print(f"raw min={raw_min:.3f}  max={raw_max:.3f}")
+    if raw_min >= -1.5 and raw_max <= 1.5:
+        images = (images + 1.0) / 2.0
+        images = np.clip(images, 0.0, 1.0)
+        return [images[i] for i in range(images.shape[0])]
+    else:
+        out = []
+        for i in range(images.shape[0]):
+            img = images[i].astype(np.float32)
+            img = (img - img.min()) / (img.max() - img.min() + 1e-8)
+            out.append(img)
+        return out
 
 
 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)
+    dummy_alpha = jnp.array(1.0)
+    dummy_rng   = jax.random.PRNGKey(0)
+    dummy_z     = jnp.zeros((1, LATENT_DIM))
+    for res in RESOLUTIONS:
+        generator.current_resolution = res
+        _ = generator(dummy_z, alpha=dummy_alpha, rng_key=dummy_rng)
+    generator.current_resolution = TARGET_RES
+    print(f"✅ Model built — trainable: {len(generator.trainable_variables)}  non-trainable: {len(generator.non_trainable_variables)}")
 
 
 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
+    global _gen_state, _use_stateless, _saved_alpha
     _build_model()
 
-    if os.path.isfile(WEIGHTS_PATH):
-        generator.load_weights(WEIGHTS_PATH)
+    if os.path.isfile(WEIGHTS_H5):
+        print(f"📂 Loading Keras weights: {WEIGHTS_H5}")
+        generator.load_weights(WEIGHTS_H5)
         _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.
-    """
+        print("✅ Keras weights loaded.")
+
+    elif os.path.isfile(CHECKPOINT_PKL):
+        print(f"📂 Loading checkpoint: {CHECKPOINT_PKL}")
+        import joblib
+        data = joblib.load(CHECKPOINT_PKL)
+        print(f"   Keys: {list(data.keys())}")
+
+        gen_trainable     = data["gen_trainable"]
+        gen_non_trainable = data.get("gen_non_trainable", [])
+        _saved_alpha      = float(data.get("alpha", 1.0))
+
+        print(f"   trainable={len(gen_trainable)}  non_trainable={len(gen_non_trainable)}  alpha={_saved_alpha:.4f}")
+
+        n_model = len(generator.trainable_variables)
+        if len(gen_trainable) != n_model:
+            raise ValueError(f"❌ Mismatch — checkpoint:{len(gen_trainable)} model:{n_model}")
+
+        n_non = len(generator.non_trainable_variables)
+        _gen_state = {
+            "ema_trainable": [jnp.asarray(t) for t in gen_trainable],
+            "non_trainable": (
+                [jnp.asarray(t) for t in gen_non_trainable[:n_non]]
+                if gen_non_trainable
+                else [jnp.asarray(v.value) for v in generator.non_trainable_variables]
+            ),
+        }
+        _use_stateless = True
+        print("✅ Checkpoint loaded.")
+
+    else:
+        raise FileNotFoundError(f"❌ No weights found in: {WEIGHTS_DIR}")
+
+    generator.current_resolution = TARGET_RES
+    print(f"🖼️  Resolution locked to: {TARGET_RES}")
+    return True
+
+
+def generate_images(n_images, resolution=None, seed=None):
+    if not _use_stateless and _gen_state is None:
+        raise RuntimeError("❌ Call load_weights() first.")
     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)
+        return _generate_stateless(n_images, rng)
+    return _generate_keras(n_images, 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
-
+def _generate_keras(n_images, rng):
+    generator.current_resolution = TARGET_RES
     rng, z_key, noise_key = jax.random.split(rng, 3)
-    z = jax.random.normal(z_key, (n_images, LATENT_DIM))
+    z   = jax.random.normal(z_key, (n_images, LATENT_DIM))
+    out = generator(z, alpha=1.0, rng_key=noise_key)
+    images = out[0] if isinstance(out, (list, tuple)) else out
+    return _postprocess(np.array(images))
 
-    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
 
+def _generate_stateless(n_images, rng):
+    generator.current_resolution = TARGET_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"],
+    images, _ = generator.stateless_call(
+        _gen_state["ema_trainable"],
+        _gen_state["non_trainable"],
         z,
-        1.0,
+        jnp.array(_saved_alpha),
         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)]
+    return _postprocess(np.array(images))

requirements.txt

@@ -4,3 +4,4 @@ keras>=3.0.0
 jax>=0.4.0
 jaxlib>=0.4.0
 numpy>=1.24.0,<2.0.0
+joblib>=1.0.0