Update app.py

app.py

@@ -57,7 +57,7 @@ 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((128, 128))  # Further reduced image size
+        image = image.convert("RGB").resize((512, 512))  # Keep original size
         image = np.array(image).astype(np.float32) / 255.0
         return image.transpose(2, 0, 1)
 
@@ -97,6 +97,7 @@ def clear_jit_cache():
 # Training function
 def train_step(state, batch, rng):
     def compute_loss(params, pixel_values, rng):
+        # Encode images to latent space
         latents = pipeline.vae.apply(
             {"params": params["vae"]},
             pixel_values,
@@ -104,10 +105,15 @@ def train_step(state, batch, rng):
         ).latent_dist.sample(rng)
         latents = latents * 0.18215
 
+        # Generate random noise
         noise = jax.random.normal(rng, latents.shape)
+        
+        # Sample random timesteps
         timesteps = jax.random.randint(
             rng, (latents.shape[0],), 0, pipeline.scheduler.config.num_train_timesteps
         )
+        
+        # Add noise to latents
         noisy_latents = pipeline.scheduler.add_noise(
             pipeline.scheduler.create_state(),
             original_samples=latents,
@@ -115,11 +121,13 @@ def train_step(state, batch, rng):
             timesteps=timesteps
         )
         
+        # Generate random encoder hidden states (simulating text embeddings)
         encoder_hidden_states = jax.random.normal(
             rng, 
             (latents.shape[0], pipeline.text_encoder.config.hidden_size)
         )
         
+        # Predict noise
         model_output = state.apply_fn.apply(
             {'params': params["unet"]},
             noisy_latents,
@@ -128,6 +136,7 @@ def train_step(state, batch, rng):
             train=True,
         )
         
+        # Compute loss
         return jnp.mean((model_output - noise) ** 2)
 
     grad_fn = jax.value_and_grad(compute_loss)
@@ -136,18 +145,9 @@ def train_step(state, batch, rng):
     state = state.apply_gradients(grads=grads)
     return state, loss
 
-# Initialize training state
-learning_rate = 1e-5
-optimizer = optax.adam(learning_rate)
-state = train_state.TrainState.create(
-    apply_fn=unet,
-    params={"unet": params["unet"], "vae": params["vae"]},
-    tx=optimizer,
-)
-
 # Training loop
-num_epochs = 3  # Further reduced number of epochs
-batch_size = 2  # Reduced batch size for CPU
+num_epochs = 3
+batch_size = 1  # Reduced batch size due to memory constraints
 rng = jax.random.PRNGKey(0)
 
 for epoch in range(num_epochs):