[Link]
ai/public/artifacts/ccfbe489-b412-42ec-984e-ac88b7c79f39
[Link]
Google collab
# Install these in Google Colab
!pip install torch torchvision
!pip install diffusers transformers
!pip install accelerate
# STEP 1: Install and Import Libraries
import torch
import [Link] as nn
from torchvision import datasets, transforms
from [Link] import DataLoader
import [Link] as plt
# Check if GPU is available
device = 'cuda' if [Link].is_available() else 'cpu'
print(f"Using: {device}")
# STEP 2: Settings (Don't worry about these details yet)
latent_dim = 100 # Size of random input
image_size = 28 # Size of output image (28x28)
batch_size = 64 # How many images to process at once
lr = 0.0002 # Learning rate (how fast it learns)
epochs = 20 # How many times to see all data
# STEP 3: Create Generator (The Artist)
class Generator([Link]):
def __init__(self):
super().__init__()
[Link] = [Link](
[Link](100, 256), # Input layer
[Link](), # Activation
[Link](256, 512), # Hidden layer
[Link](),
[Link](512, 784), # Output (28*28=784)
[Link]() # Make output between -1 and 1
)
def forward(self, z):
return [Link](z).view(-1, 1, 28, 28)
# STEP 4: Create Discriminator (The Critic)
class Discriminator([Link]):
def __init__(self):
super().__init__()
[Link] = [Link](
[Link](784, 512), # Input
[Link](0.2),
[Link](512, 256), # Hidden layer
[Link](0.2),
[Link](256, 1), # Output: Real or Fake?
[Link]() # Output between 0 and 1
)
def forward(self, img):
return [Link]([Link](-1, 784))
# STEP 5: Initialize Both Models
generator = Generator().to(device)
discriminator = Discriminator().to(device)
# STEP 6: Set Up Training Tools
criterion = [Link]() # Loss function
opt_g = [Link]([Link](), lr=lr)
opt_d = [Link]([Link](), lr=lr)
# STEP 7: Load Training Data (MNIST digits)
transform = [Link]([
[Link](),
[Link]([0.5], [0.5]) # Normalize to [-1, 1]
])
dataset = [Link](root='./data', train=True, download=True, transform=transform)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
print(f"Dataset loaded: {len(dataset)} images")
# STEP 8: Training Loop
print("Starting training...")
for epoch in range(epochs):
for i, (real_images, _) in enumerate(dataloader):
real_images = real_images.to(device)
batch_size_current = real_images.size(0)
# Labels
real_labels = [Link](batch_size_current, 1).to(device)
fake_labels = [Link](batch_size_current, 1).to(device)
# ============ Train Discriminator ============
# On real images
outputs = discriminator(real_images)
d_loss_real = criterion(outputs, real_labels)
# On fake images
z = [Link](batch_size_current, latent_dim).to(device)
fake_images = generator(z)
outputs = discriminator(fake_images.detach())
d_loss_fake = criterion(outputs, fake_labels)
# Total discriminator loss
d_loss = d_loss_real + d_loss_fake
opt_d.zero_grad()
d_loss.backward()
opt_d.step()
# ============ Train Generator ============
z = [Link](batch_size_current, latent_dim).to(device)
fake_images = generator(z)
outputs = discriminator(fake_images)
g_loss = criterion(outputs, real_labels) # Want discriminator to think these are real
opt_g.zero_grad()
g_loss.backward()
opt_g.step()
# Print progress every 100 batches
if i % 100 == 0:
print(f'Epoch [{epoch}/{epochs}] Batch [{i}/{len(dataloader)}] '
f'D_loss: {d_loss.item():.4f} G_loss: {g_loss.item():.4f}')
# Show generated images every 5 epochs
if epoch % 5 == 0:
with torch.no_grad():
z = [Link](16, latent_dim).to(device)
fake = generator(z).cpu()
fig, axes = [Link](4, 4, figsize=(8,8))
for idx, ax in enumerate([Link]):
img = fake[idx].squeeze()
img = (img + 1) / 2 # Denormalize
[Link](img, cmap='gray')
[Link]('off')
[Link](f'Epoch {epoch}')
[Link]()
print("Training complete!")
# STEP 9: Generate Final Images
with torch.no_grad():
z = [Link](25, latent_dim).to(device)
generated = generator(z).cpu()
fig, axes = [Link](5, 5, figsize=(10,10))
for idx, ax in enumerate([Link]):
img = generated[idx].squeeze()
img = (img + 1) / 2
[Link](img, cmap='gray')
[Link]('off')
[Link]('Final Generated Digits')
[Link]()
# STEP 10: Save the Model
[Link](generator.state_dict(), '[Link]')
print("Model saved!")
Using Gradio
# Install gradio
!pip install gradio
import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
# Load your model
pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16
).to("cuda")
# Load your LoRA
pipe.load_lora_weights("./my_model")
# Create generation function
def generate_image(prompt):
image = pipe(prompt, num_inference_steps=50).images[0]
return image
# Create web interface
demo = [Link](
fn=generate_image,
inputs=[Link](label="Enter your prompt", placeholder="a photo of sks dog on the
beach"),
outputs=[Link](label="Generated Image"),
title="My Custom Image Generator",
description="Generate images using my trained model!"
)
# Launch
[Link](share=True) # share=True creates public link
Creates a web interface
Gives you a public link
Anyone can visit and generate images
Runs for 72 hours (in Colab)