Update app.py

app.py

@@ -12,6 +12,7 @@ print(f'Using device: {device}')
 # Create a scheduler
 noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule='squaredcos_cap_v2')
 
+
 class ClassConditionedUnet(nn.Module):
   def __init__(self, num_classes=10, class_emb_size=4):
     super().__init__()
@@ -58,12 +59,6 @@ class ClassConditionedUnet(nn.Module):
     return self.model(net_input, t).sample # (bs, 1, 28, 28)
 
 
-# CIFAR-10 class names
-cifar10_classes = [
-    "plane", "car", "bird", "cat", "deer",
-    "dog", "frog", "horse", "ship", "truck"
-]
-
 
 def load_checkpoint_for_inference(filepath, model_class):
     """
@@ -95,9 +90,6 @@ def load_checkpoint_for_inference(filepath, model_class):
     # loaded_model = load_checkpoint_for_inference("model_path", ClassConditionedUnet)
 
 
-# Initialize a dummy model (replace with your actual model loading)
-model = load_checkpoint_for_inference(filepath="/content/drive/MyDrive/Colab Notebooks/HF_Diffusion_Course/model_v02/CIFAR10_unet_v_02_100_epochs_inference.pth", model_class=ClassConditionedUnet)
-
 
 def generate_images(selected_class_name, num_samples=4):
     print(f"Generating {num_samples} samples for class: {selected_class_name}")
@@ -138,6 +130,18 @@ def generate_images(selected_class_name, num_samples=4):
 
 
 
+
+
+# Initialize a dummy model (replace with your actual model loading)
+model = load_checkpoint_for_inference(filepath="CIFAR10_unet_v_02_100_epochs_inference.pth", model_class=ClassConditionedUnet)
+
+# CIFAR-10 class names
+cifar10_classes = [
+    "plane", "car", "bird", "cat", "deer",
+    "dog", "frog", "horse", "ship", "truck"
+]
+
+
 # Create the Gradio interface
 custom_css = """
 #gallery {