Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| import os | |
| import legacy | |
| import torch_utils | |
| # Load the pre-trained StyleGAN model | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| model_path = 'dress_model.pkl' # Place your .pkl in the same directory or update path | |
| # Load StyleGAN Generator | |
| with open(model_path, 'rb') as f: | |
| G = legacy.load_network_pkl(f)['G_ema'].to(device) | |
| def mix_styles(image1_path, image2_path, styles_to_mix): | |
| # Extract image names (without extensions) | |
| image1_name = os.path.splitext(os.path.basename(image1_path))[0] | |
| image2_name = os.path.splitext(os.path.basename(image2_path))[0] | |
| # Load latent vectors from .npz | |
| latent_vector_1 = np.load(os.path.join("projection_results", image1_name, "projected_w.npz"))['w'] | |
| latent_vector_2 = np.load(os.path.join("projection_results", image2_name, "projected_w.npz"))['w'] | |
| # Convert to torch tensors | |
| latent_1_tensor = torch.from_numpy(latent_vector_1).to(device) | |
| latent_2_tensor = torch.from_numpy(latent_vector_2).to(device) | |
| # Mix layers | |
| mixed_latent = latent_1_tensor.clone() | |
| mixed_latent[:, styles_to_mix] = latent_2_tensor[:, styles_to_mix] | |
| # Generate image | |
| with torch.no_grad(): | |
| image = G.synthesis(mixed_latent, noise_mode='const') | |
| # Convert to image | |
| image = (image.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).cpu().numpy() | |
| mixed_image = Image.fromarray(image[0], 'RGB') | |
| return mixed_image | |
| def style_mixing_interface(image1, image2, mix_value): | |
| if image1 is None or image2 is None: | |
| return None | |
| selected_layers = list(range(mix_value + 1)) | |
| return mix_styles(image1, image2, selected_layers) | |
| # Gradio UI | |
| iface = gr.Interface( | |
| fn=style_mixing_interface, | |
| inputs=[ | |
| gr.Image(label="First Clothing Image", type="filepath"), | |
| gr.Image(label="Second Clothing Image", type="filepath"), | |
| gr.Slider(label="Style Mixing Strength (Layers 0 to N)", minimum=0, maximum=9, step=1, value=5) | |
| ], | |
| outputs=gr.Image(label="Mixed Clothing Design"), | |
| title="Style Mixing for Clothing Design", | |
| description="Upload two projected images and choose how many early layers to mix. Precomputed latent vectors (projected_w.npz) must be in 'projection_results/{image_name}/'." | |
| ) | |
| iface.launch() | |