Spaces:
Build error
Build error
File size: 2,258 Bytes
da0313f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import torch
import torchvision.transforms as transforms
from PIL import Image
import gradio as gr
from tqdm import tqdm
def optimize_latent_vector(G, target_image, num_iterations=1000):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
target_image = transforms.Resize((G.img_resolution, G.img_resolution))(target_image)
target_tensor = transforms.ToTensor()(target_image).unsqueeze(0).to(device)
target_tensor = (target_tensor * 2) - 1 # Normalize to [-1, 1]
latent_vector = torch.randn((1, G.z_dim), device=device, requires_grad=True)
optimizer = torch.optim.Adam([latent_vector], lr=0.1)
for i in tqdm(range(num_iterations), desc="Optimizing latent vector"):
optimizer.zero_grad()
generated_image = G(latent_vector, None)
loss = torch.nn.functional.mse_loss(generated_image, target_tensor)
loss.backward()
optimizer.step()
if (i + 1) % 100 == 0:
print(f'Iteration {i+1}/{num_iterations}, Loss: {loss.item()}')
return latent_vector.detach()
def generate_from_upload(uploaded_image):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Optimize latent vector for the uploaded image
optimized_z = optimize_latent_vector(G, uploaded_image)
# Generate variations
num_variations = 4
variation_strength = 0.1
varied_z = optimized_z + torch.randn((num_variations, G.z_dim), device=device) * variation_strength
# Generate the variations
with torch.no_grad():
imgs = G(varied_z, c=None, truncation_psi=0.7, noise_mode='const')
imgs = (imgs * 127.5 + 128).clamp(0, 255).to(torch.uint8)
imgs = imgs.permute(0, 2, 3, 1).cpu().numpy()
# Convert the generated image tensors to PIL Images
generated_images = [Image.fromarray(img) for img in imgs]
# Return the images separately
return generated_images[0], generated_images[1], generated_images[2], generated_images[3]
# Create the Gradio interface
iface = gr.Interface(
fn=generate_from_upload,
inputs=gr.Image(type="pil"),
outputs=[gr.Image(type="pil") for _ in range(4)],
title="StyleGAN Image Variation Generator"
)
# Launch the Gradio interface
iface.launch(share=True, debug=True)
|