Question: Generate image using SDS

import math
from tqdm import tqdm
import torch
import torch.nn as nn
from nerf.sd import StableDiffusion, seed_everything
from torch.optim.lr_scheduler import LambdaLR

import matplotlib.pyplot as plt


def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, num_cycles: float = 0.5):

    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))

    return LambdaLR(optimizer, lr_lambda, -1)

device = 'cuda:0'
guidance = StableDiffusion(device)
# limited memory here, don't need the decoder
guidance.vae.decoder = None

prompt = '3D texture of pebbles'
text_embeddings = guidance.get_text_embeds(prompt, '')
guidance.text_encoder.to('cpu')
torch.cuda.empty_cache()

seed_everything(42)
# put parameters approximately in range(0, 1) since this is what `encode_imgs` expects
rgb = nn.Parameter(torch.randn(1, 3, 512, 512, device=device) / 2  + .5)
optimizer = torch.optim.AdamW([rgb], lr=1e-1, weight_decay=0)
num_steps = 5000
scheduler = get_cosine_schedule_with_warmup(optimizer, 100, int(num_steps*1.5))

for step in tqdm(range(num_steps)):
    optimizer.zero_grad()
    guidance.train_step(text_embeddings, rgb, guidance_scale=100)

    optimizer.step()
    scheduler.step()

plt.imshow(rgb.detach().clamp(0, 1).squeeze(0).permute(1,2,0).cpu())
plt.axis('off')
plt.show()