Create model.py

model.py

@@ -0,0 +1,103 @@
+import numpy as np
+import tensorflow as tf
+import matplotlib.pyplot as plt
+
+# Parameters
+IMG_SIZE = 64
+LATENT_DIM = 128  # Dimension of the input noise
+BATCH_SIZE = 1
+STYLE_DIM = 3  # Number of styles: circles, squares, mixed
+
+# Generator with style control
+def build_generator():
+    # Input layers
+    noise_input = tf.keras.layers.Input(shape=(LATENT_DIM,))
+    style_input = tf.keras.layers.Input(shape=(STYLE_DIM,))
+    
+    # Concatenate noise and style
+    x = tf.keras.layers.concatenate([noise_input, style_input])
+    x = tf.keras.layers.Dense(8 * 8 * 64, activation='relu')(x)
+    x = tf.keras.layers.Reshape((8, 8, 64))(x)
+    
+    # Transposed convolutions
+    x = tf.keras.layers.Conv2DTranspose(64, (4,4), strides=2, padding='same', activation='relu')(x)
+    x = tf.keras.layers.Conv2DTranspose(32, (4,4), strides=2, padding='same', activation='relu')(x)
+    x = tf.keras.layers.Conv2DTranspose(3, (4,4), strides=2, padding='same', activation='sigmoid')(x)  # RGB [0,1]
+    
+    return tf.keras.Model(inputs=[noise_input, style_input], outputs=x)
+
+# Discriminator (same as before)
+def build_discriminator():
+    model = tf.keras.Sequential([
+        tf.keras.layers.Conv2D(32, (3,3), strides=2, padding='same', input_shape=(IMG_SIZE, IMG_SIZE, 3)),
+        tf.keras.layers.LeakyReLU(0.2),
+        tf.keras.layers.Conv2D(64, (3,3), strides=2, padding='same'),
+        tf.keras.layers.LeakyReLU(0.2),
+        tf.keras.layers.Flatten(),
+        tf.keras.layers.Dense(1, activation='sigmoid')
+    ])
+    return model
+
+# Create models
+generator = build_generator()
+discriminator = build_discriminator()
+
+# Loss function and optimizers
+cross_entropy = tf.keras.losses.BinaryCrossentropy()
+g_optimizer = tf.keras.optimizers.Adam(0.0002, beta_1=0.5)
+d_optimizer = tf.keras.optimizers.Adam(0.0002, beta_1=0.5)
+
+# Style encodings
+STYLES = {
+    'circles': [1., 0., 0.],
+    'squares': [0., 1., 0.],
+    'mixed': [0., 0., 1.]
+}
+
+def generate_with_style(style_name):
+    style = STYLES[style_name]
+    test_noise = tf.random.normal([1, LATENT_DIM])
+    style_input = tf.constant([style], dtype=tf.float32)
+    generated_img = generator([test_noise, style_input], training=False)[0]
+    plt.imshow(generated_img)
+    plt.title(f"Style: {style_name}")
+    plt.axis('off')
+    plt.show()
+
+# Training function with style conditioning
+@tf.function
+def train_step():
+    # 1. Generate random noise and random style
+    noise = tf.random.normal([BATCH_SIZE, LATENT_DIM])
+    style = tf.one_hot(tf.random.uniform([BATCH_SIZE], maxval=STYLE_DIM, dtype=tf.int32), STYLE_DIM)
+    
+    # 2. Train discriminator
+    with tf.GradientTape() as d_tape:
+        generated_images = generator([noise, style], training=True)
+        real_output = discriminator(tf.random.uniform((BATCH_SIZE, IMG_SIZE, IMG_SIZE, 3)), training=True)
+        fake_output = discriminator(generated_images, training=True)
+        d_loss = cross_entropy(tf.ones_like(fake_output), fake_output)
+    
+    d_gradients = d_tape.gradient(d_loss, discriminator.trainable_variables)
+    d_optimizer.apply_gradients(zip(d_gradients, discriminator.trainable_variables))
+    
+    # 3. Train generator
+    with tf.GradientTape() as g_tape:
+        generated_images = generator([noise, style], training=True)
+        fake_output = discriminator(generated_images, training=True)
+        g_loss = cross_entropy(tf.ones_like(fake_output), fake_output)
+    
+    g_gradients = g_tape.gradient(g_loss, generator.trainable_variables)
+    g_optimizer.apply_gradients(zip(g_gradients, generator.trainable_variables))
+    
+    return d_loss, g_loss
+
+# Training loop
+for epoch in range(50):
+    d_loss, g_loss = train_step()
+    if epoch % 10 == 0:
+        print(f"Epoch {epoch}, D Loss: {d_loss:.3f}, G Loss: {g_loss:.3f}")
+        
+        # Generate samples for each style
+        for style_name in STYLES.keys():
+            generate_with_style(style_name)