Update app.py

app.py

@@ -1,6 +1,5 @@
 import jax
 import jax.numpy as jnp
-from flax.jax_utils import replicate
 from flax.training import train_state
 import optax
 from diffusers import FlaxStableDiffusionPipeline
@@ -12,6 +11,11 @@ from PIL import Image
 import numpy as np
 import gc
 
+# Force JAX to use CPU
+jax.config.update('jax_platform_name', 'cpu')
+
+print("Using CPU for computations")
+
 # Set up cache directories
 cache_dir = "/tmp/huggingface_cache"
 model_cache_dir = os.path.join(cache_dir, "stable_diffusion_model")
@@ -33,7 +37,7 @@ def get_model(model_id, revision):
         pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
             model_id, 
             revision=revision,
-            dtype=jnp.float16,
+            dtype=jnp.float32,  # Use float32 for CPU
         )
         with open(model_cache_file, 'wb') as f:
             pickle.dump((pipeline, params), f)
@@ -53,8 +57,8 @@ def preprocess_images(examples):
             image = Image.open(image)
         if not isinstance(image, Image.Image):
             raise ValueError(f"Unexpected image type: {type(image)}")
-        image = image.convert("RGB").resize((256, 256))  # Reduced image size
-        image = np.array(image).astype(np.float16) / 255.0
+        image = image.convert("RGB").resize((128, 128))  # Further reduced image size
+        image = np.array(image).astype(np.float32) / 255.0
         return image.transpose(2, 0, 1)
 
     return {"pixel_values": [process_image(img) for img in examples["image"]]}
@@ -73,7 +77,7 @@ try:
             processed_dataset = pickle.load(f)
     else:
         print("Loading dataset from Hugging Face...")
-        dataset = load_dataset(dataset_name, split="train[:1000]")  # Load only first 1000 samples
+        dataset = load_dataset(dataset_name, split="train[:500]")  # Load only first 500 samples
         print("Processing dataset...")
         processed_dataset = dataset.map(preprocess_images, batched=True, remove_columns=dataset.column_names)
         with open(dataset_cache_file, 'wb') as f:
@@ -90,8 +94,7 @@ def clear_jit_cache():
     jax.clear_caches()
     gc.collect()
 
-# Training function with gradient accumulation
-@jax.jit
+# Training function
 def train_step(state, batch, rng):
     def compute_loss(params, pixel_values, rng):
         latents = pipeline.vae.apply(
@@ -99,9 +102,9 @@ def train_step(state, batch, rng):
             pixel_values,
             method=pipeline.vae.encode
         ).latent_dist.sample(rng)
-        latents = latents * jnp.float16(0.18215)
+        latents = latents * 0.18215
 
-        noise = jax.random.normal(rng, latents.shape, dtype=jnp.float16)
+        noise = jax.random.normal(rng, latents.shape)
         timesteps = jax.random.randint(
             rng, (latents.shape[0],), 0, pipeline.scheduler.config.num_train_timesteps
         )
@@ -114,8 +117,7 @@ def train_step(state, batch, rng):
         
         encoder_hidden_states = jax.random.normal(
             rng, 
-            (latents.shape[0], pipeline.text_encoder.config.hidden_size),
-            dtype=jnp.float16
+            (latents.shape[0], pipeline.text_encoder.config.hidden_size)
         )
         
         model_output = state.apply_fn.apply(
@@ -135,7 +137,7 @@ def train_step(state, batch, rng):
     return state, loss
 
 # Initialize training state
-learning_rate = jnp.float16(1e-5)
+learning_rate = 1e-5
 optimizer = optax.adam(learning_rate)
 state = train_state.TrainState.create(
     apply_fn=unet,
@@ -144,15 +146,15 @@ state = train_state.TrainState.create(
 )
 
 # Training loop
-num_epochs = 5  # Reduced number of epochs
-batch_size = 4
+num_epochs = 3  # Further reduced number of epochs
+batch_size = 2  # Reduced batch size for CPU
 rng = jax.random.PRNGKey(0)
 
 for epoch in range(num_epochs):
     epoch_loss = 0
     num_batches = 0
     for batch in tqdm(processed_dataset.batch(batch_size)):
-        batch['pixel_values'] = jnp.array(batch['pixel_values'], dtype=jnp.float16)
+        batch['pixel_values'] = jnp.array(batch['pixel_values'])
         rng, step_rng = jax.random.split(rng)
         state, loss = train_step(state, batch, step_rng)
         epoch_loss += loss