Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import torch.nn as nn | |
| import torchvision.transforms as transforms | |
| import torchvision.models as models | |
| from PIL import Image | |
| import numpy as np | |
| import os | |
| import requests | |
| import base64 | |
| import io | |
| # CPU optimization: Disable CUDA and use optimized CPU threads | |
| torch.set_num_threads(4) # Adjust based on your CPU | |
| device = torch.device("cpu") | |
| # Get QuickCloud API URL from environment variable | |
| QUICKCLOUD_API_URL = os.environ.get("QUICKCLOUD_API_URL", "") | |
| class LightingStyleTransfer: | |
| def __init__(self): | |
| # Use VGG16 for feature extraction (lighter than VGG19) | |
| vgg = models.vgg16(pretrained=True).features.to(device).eval() | |
| # Freeze parameters for CPU efficiency | |
| for param in vgg.parameters(): | |
| param.requires_grad = False | |
| self.model = vgg | |
| # Layer indices for content and style | |
| self.style_layers = [0, 5, 10, 17] # Reduced layers for CPU | |
| self.content_layers = [17] | |
| def preprocess(self, img, max_size=512): | |
| """Resize and normalize image - smaller size for CPU""" | |
| # CPU optimization: Use smaller image size | |
| w, h = img.size | |
| scale = max_size / max(w, h) | |
| new_size = (int(w * scale), int(h * scale)) | |
| img = img.resize(new_size, Image.LANCZOS) | |
| transform = transforms.Compose([ | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]) | |
| ]) | |
| return transform(img).unsqueeze(0).to(device) | |
| def deprocess(self, tensor): | |
| """Convert tensor back to image""" | |
| img = tensor.cpu().clone().squeeze(0) | |
| img = img.clamp(0, 1) | |
| img = transforms.ToPILImage()(img) | |
| return img | |
| def gram_matrix(self, tensor): | |
| """Compute Gram matrix for style representation""" | |
| b, c, h, w = tensor.size() | |
| features = tensor.view(b * c, h * w) | |
| G = torch.mm(features, features.t()) | |
| return G.div(b * c * h * w) | |
| def get_features(self, image): | |
| """Extract features from specified layers""" | |
| features = {} | |
| x = image | |
| for idx, layer in enumerate(self.model): | |
| x = layer(x) | |
| if idx in self.style_layers: | |
| features[f'style_{idx}'] = x | |
| if idx in self.content_layers: | |
| features[f'content_{idx}'] = x | |
| return features | |
| def transfer(self, content_img, style_img, steps=150, style_weight=1e6, | |
| content_weight=1): | |
| """Perform lighting style transfer""" | |
| # Preprocess images | |
| content = self.preprocess(content_img) | |
| style = self.preprocess(style_img) | |
| # Initialize target as content image | |
| target = content.clone().requires_grad_(True) | |
| # Get features | |
| content_features = self.get_features(content) | |
| style_features = self.get_features(style) | |
| # Compute style gram matrices | |
| style_grams = {k: self.gram_matrix(v) for k, v in style_features.items() | |
| if 'style' in k} | |
| # CPU optimization: Use LBFGS optimizer (faster convergence) | |
| optimizer = torch.optim.LBFGS([target], max_iter=20) | |
| step = [0] | |
| def closure(): | |
| target.data.clamp_(0, 1) | |
| optimizer.zero_grad() | |
| target_features = self.get_features(target) | |
| # Content loss | |
| content_loss = 0 | |
| for k in content_features: | |
| if 'content' in k: | |
| content_loss += torch.mean((target_features[k] - | |
| content_features[k]) ** 2) | |
| # Style loss | |
| style_loss = 0 | |
| for k in style_grams: | |
| target_gram = self.gram_matrix(target_features[k]) | |
| style_loss += torch.mean((target_gram - style_grams[k]) ** 2) | |
| # Total loss | |
| total_loss = content_weight * content_loss + style_weight * style_loss | |
| total_loss.backward() | |
| step[0] += 1 | |
| if step[0] % 30 == 0: | |
| print(f"Step {step[0]}, Loss: {total_loss.item():.2f}") | |
| return total_loss | |
| # Optimization loop | |
| epochs = steps // 20 # LBFGS takes ~20 iterations per step | |
| for i in range(epochs): | |
| optimizer.step(closure) | |
| if step[0] >= steps: | |
| break | |
| # Final clamp and return | |
| target.data.clamp_(0, 1) | |
| return self.deprocess(target) | |
| def process_with_quickcloud(content_img, style_img, steps, style_strength): | |
| """Process using QuickCloud API (powered by Modal.com)""" | |
| if not QUICKCLOUD_API_URL: | |
| return None, "❌ QuickCloud API URL not configured. Please set QUICKCLOUD_API_URL environment variable." | |
| try: | |
| # Convert PIL images to bytes | |
| content_bytes = io.BytesIO() | |
| style_bytes = io.BytesIO() | |
| content_img.save(content_bytes, format='PNG') | |
| style_img.save(style_bytes, format='PNG') | |
| # Encode to base64 | |
| content_b64 = base64.b64encode(content_bytes.getvalue()).decode() | |
| style_b64 = base64.b64encode(style_bytes.getvalue()).decode() | |
| # Prepare request | |
| payload = { | |
| "content_image": content_b64, | |
| "style_image": style_b64, | |
| "steps": steps, | |
| "style_weight": style_strength * 1e6, | |
| "content_weight": 1.0, | |
| "learning_rate": 0.03 | |
| } | |
| print("Sending request to NamelessAI QuickCloud (H100 GPU)...") | |
| # Make API request | |
| response = requests.post(QUICKCLOUD_API_URL, json=payload, timeout=300) | |
| response.raise_for_status() | |
| # Decode result | |
| result_data = response.json() | |
| result_bytes = base64.b64decode(result_data["result_image"]) | |
| result_img = Image.open(io.BytesIO(result_bytes)) | |
| return result_img, "✅ Processing complete via QuickCloud (H100 GPU)!" | |
| except requests.exceptions.Timeout: | |
| return None, "❌ Request timed out. Please try again." | |
| except requests.exceptions.RequestException as e: | |
| return None, f"❌ API Error: {str(e)}" | |
| except Exception as e: | |
| return None, f"❌ Error: {str(e)}" | |
| def process_locally(content_img, style_img, steps, style_strength): | |
| """Process using local CPU""" | |
| try: | |
| # Adjust style weight | |
| style_weight = style_strength * 1e6 | |
| # Perform transfer | |
| result = style_transfer.transfer( | |
| content_img, | |
| style_img, | |
| steps=steps, | |
| style_weight=style_weight, | |
| content_weight=1 | |
| ) | |
| return result, "✅ Processing complete via Local CPU!" | |
| except Exception as e: | |
| return None, f"❌ Error: {str(e)}" | |
| def process_images(content_img, style_img, steps, style_strength, use_quickcloud): | |
| """Process the style transfer based on selected mode""" | |
| if content_img is None or style_img is None: | |
| return None, "⚠️ Please upload both content and style images." | |
| # Convert to PIL if needed | |
| if isinstance(content_img, np.ndarray): | |
| content_img = Image.fromarray(content_img) | |
| if isinstance(style_img, np.ndarray): | |
| style_img = Image.fromarray(style_img) | |
| if use_quickcloud: | |
| return process_with_quickcloud(content_img, style_img, steps, style_strength) | |
| else: | |
| return process_locally(content_img, style_img, steps, style_strength) | |
| # Initialize local model (done once at startup) | |
| print("Loading local model... This may take a moment.") | |
| style_transfer = LightingStyleTransfer() | |
| print("Local model loaded successfully!") | |
| # Check if QuickCloud is available | |
| quickcloud_available = bool(QUICKCLOUD_API_URL) | |
| if quickcloud_available: | |
| print(f"✓ QuickCloud API configured and available") | |
| else: | |
| print("✗ QuickCloud API not configured (set QUICKCLOUD_API_URL environment variable)") | |
| # Create Gradio interface | |
| with gr.Blocks(title="AI Lighting Style Transfer") as demo: | |
| gr.Markdown(""" | |
| # 🎨 AI-Powered Lighting Style Transfer | |
| Transfer the lighting and color style from one image to another using neural style transfer. | |
| ## Processing Options: | |
| - **Local (CPU)**: Runs on your machine. Takes 1-3 minutes. Free. | |
| - **NamelessAI QuickCloud**: Runs on H100 GPU cloud. Takes 5-10 seconds. Requires API key. | |
| - *Powered by Modal.com* | |
| ## How to use: | |
| 1. Upload your **content image** (the image you want to transform) | |
| 2. Upload your **style image** (the image whose lighting you want to copy) | |
| 3. Choose processing mode (Local or QuickCloud) | |
| 4. Adjust settings if desired | |
| 5. Click "Transfer Style" and wait for processing | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| content_input = gr.Image(label="Content Image", type="pil") | |
| style_input = gr.Image(label="Style Image", type="pil") | |
| with gr.Column(): | |
| output = gr.Image(label="Result") | |
| status_text = gr.Textbox(label="Status", interactive=False) | |
| with gr.Row(): | |
| use_quickcloud = gr.Checkbox( | |
| label="Use NamelessAI QuickCloud (H100 GPU - Powered by Modal.com)", | |
| value=False, | |
| interactive=quickcloud_available, | |
| info="5-10 seconds vs 1-3 minutes locally" if quickcloud_available else "API URL not configured" | |
| ) | |
| with gr.Row(): | |
| steps_slider = gr.Slider( | |
| minimum=50, | |
| maximum=300, | |
| value=150, | |
| step=10, | |
| label="Optimization Steps (more = better quality, slower)" | |
| ) | |
| style_strength = gr.Slider( | |
| minimum=0.5, | |
| maximum=3.0, | |
| value=1.0, | |
| step=0.1, | |
| label="Style Strength" | |
| ) | |
| transfer_btn = gr.Button("Transfer Style", variant="primary", size="lg") | |
| gr.Markdown(""" | |
| ### Tips: | |
| - **Local Mode**: Images resized to 512px, use 100-150 steps for balance | |
| - **QuickCloud Mode**: Handles 1024px images, 300 steps recommended for best quality | |
| - Increase style strength for more dramatic lighting effects | |
| - Works best with images that have distinct lighting patterns | |
| ### QuickCloud Setup: | |
| To use QuickCloud, set the `QUICKCLOUD_API_URL` environment variable to your Modal API endpoint. | |
| """) | |
| # Set up the button click | |
| transfer_btn.click( | |
| fn=process_images, | |
| inputs=[content_input, style_input, steps_slider, style_strength, use_quickcloud], | |
| outputs=[output, status_text] | |
| ) | |
| # Launch the app | |
| if __name__ == "__main__": | |
| demo.launch() |