Spaces:
Running
Running
| import gradio as gr | |
| import requests | |
| import time | |
| import io | |
| import os | |
| from PIL import Image | |
| # Global variable to track active requests | |
| active_requests = {} | |
| # Predefined assets | |
| ASSETS = { | |
| "top": { | |
| "reference_image": "pullover.png", | |
| "prompt": "black pullover with the logo on the chest" | |
| }, | |
| "bottom": { | |
| "reference_image": "sweatpants.png", | |
| "prompt": "black sweatpants with the silver logo on it" | |
| }, | |
| "logo": "logo.png" | |
| } | |
| def needs_resizing(image_path, max_short_edge=1280): | |
| # Check if image needs resizing (returns True if short edge > max_short_edge) | |
| with Image.open(image_path) as img: | |
| width, height = img.size | |
| return min(width, height) > max_short_edge | |
| def resize_image_to_short_edge(image_path, max_short_edge=1280): | |
| # Resize image so the short edge is at most max_short_edge pixels | |
| with Image.open(image_path) as img: | |
| width, height = img.size | |
| format = img.format | |
| # Determine scaling factor | |
| if width < height: | |
| scaling_factor = max_short_edge / width | |
| else: | |
| scaling_factor = max_short_edge / height | |
| # Calculate new dimensions and resize | |
| new_width = int(width * scaling_factor) | |
| new_height = int(height * scaling_factor) | |
| resized_img = img.resize((new_width, new_height), Image.LANCZOS) | |
| # Save to bytes buffer | |
| buffer = io.BytesIO() | |
| resized_img.save(buffer, format=format) | |
| buffer.seek(0) | |
| return buffer | |
| def validate_assets(): | |
| # Check if all required asset files exist | |
| missing = [] | |
| for asset_type in ["top", "bottom"]: | |
| if not os.path.exists(ASSETS[asset_type]["reference_image"]): | |
| missing.append(ASSETS[asset_type]["reference_image"]) | |
| if not os.path.exists(ASSETS["logo"]): | |
| missing.append(ASSETS["logo"]) | |
| if missing: | |
| raise FileNotFoundError(f"Missing required asset files: {', '.join(missing)}") | |
| def generate_and_wait_for_image( | |
| api_key: str, | |
| input_image: str, | |
| garment_type: str, | |
| progress=gr.Progress() | |
| ): | |
| # Make POST request and automatically poll for the result image | |
| # Validate assets first | |
| try: | |
| validate_assets() | |
| except FileNotFoundError as e: | |
| return None, str(e) | |
| # Create a unique ID for this request | |
| request_id = str(time.time()) | |
| active_requests[request_id] = True | |
| try: | |
| # Start the job | |
| post_url = "https://api.stability.ai/private/alo/v1/vto-acolade" | |
| headers = { | |
| "Authorization": f"Bearer {api_key}", | |
| "Accept": "application/json" | |
| } | |
| files = {} | |
| try: | |
| progress(0, desc="Starting image generation...") | |
| # Prepare all required files | |
| files = { | |
| 'logo_image': open(ASSETS["logo"], 'rb'), | |
| 'reference_image': open(ASSETS[garment_type]["reference_image"], 'rb') | |
| } | |
| # Only resize input_image if needed | |
| if needs_resizing(input_image): | |
| files['input_image'] = resize_image_to_short_edge(input_image) | |
| else: | |
| files['input_image'] = open(input_image, 'rb') | |
| data = { | |
| 'reference_image_type': (f'{garment_type}'), | |
| 'output_format': (None, "png") # Hardcoded to PNG | |
| } | |
| # Submit the job with timeout | |
| print(headers) | |
| print(files) | |
| response = requests.post(post_url, headers=headers, files=files, data=data, timeout=10) | |
| response.raise_for_status() | |
| job_data = response.json() | |
| job_id = job_data.get('id') | |
| if not job_id: | |
| return None, "Error: No job ID received in response" | |
| # Now poll for results with optimized timing | |
| get_url = f"https://api.stability.ai/private/alo/v1/results/{job_id}" | |
| headers = { | |
| "authorization": f"{api_key}", | |
| "accept": "*/*" | |
| } | |
| progress(0.3, desc="Processing your image...") | |
| # Optimized polling strategy | |
| max_attempts = 20 | |
| initial_delay = 1.0 | |
| max_delay = 5.0 | |
| current_delay = initial_delay | |
| for attempt in range(max_attempts): | |
| if not active_requests.get(request_id, False): | |
| return None, "Request cancelled by user" | |
| time.sleep(current_delay) | |
| progress(0.3 + (0.7 * attempt/max_attempts), | |
| desc=f"Checking status (attempt {attempt + 1}/{max_attempts})") | |
| try: | |
| with requests.Session() as session: | |
| response = session.get(get_url, headers=headers, timeout=10) | |
| if response.status_code == 200: | |
| if 'image' in response.headers.get('Content-Type', ''): | |
| img = Image.open(io.BytesIO(response.content)) | |
| progress(1.0, desc="Done!") | |
| return img, f"Success! Job ID: {job_id}" | |
| else: | |
| json_response = response.json() | |
| if json_response.get('status') == 'processing': | |
| current_delay = min(current_delay * 1.5, max_delay) | |
| continue | |
| return None, f"API response: {json_response}" | |
| elif response.status_code == 202: | |
| current_delay = min(current_delay * 1.5, max_delay) | |
| continue | |
| else: | |
| response.raise_for_status() | |
| except requests.exceptions.RequestException: | |
| current_delay = min(current_delay * 1.5, max_delay) | |
| continue | |
| return None, f"Timeout after {max_attempts} attempts. Job ID: {job_id}" | |
| except Exception as e: | |
| return None, f"Error: {str(e)}" | |
| finally: | |
| # Clean up file handles | |
| for key in files: | |
| if hasattr(files[key], 'close'): | |
| files[key].close() | |
| finally: | |
| # Clean up the request tracking | |
| active_requests.pop(request_id, None) | |
| def cancel_request(): | |
| """Function to cancel active requests""" | |
| for req_id in list(active_requests.keys()): | |
| active_requests[req_id] = False | |
| with gr.Blocks(title="Virtual Try-On Demo") as demo: | |
| gr.Markdown(""" | |
| # Virtual Try-On Demo v1 | |
| Upload your photo and select garment type to generate your VTon image. | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| api_key = gr.Textbox( | |
| label="API Key", | |
| value="", | |
| type="password" | |
| ) | |
| input_image = gr.Image( | |
| label="Upload Your Photo", | |
| type="filepath", | |
| sources=["upload"], | |
| height=300 | |
| ) | |
| garment_type = gr.Dropdown( | |
| label="Garment Type", | |
| choices=["top", "bottom"], | |
| value="top" | |
| ) | |
| with gr.Row(): | |
| submit_btn = gr.Button("Generate Image", variant="primary") | |
| cancel_btn = gr.Button("Cancel Request") | |
| with gr.Column(): | |
| output_image = gr.Image( | |
| label="Generated Result", | |
| interactive=False, | |
| height=400 | |
| ) | |
| status_output = gr.Textbox( | |
| label="Status", | |
| interactive=False | |
| ) | |
| submit_btn.click( | |
| fn=generate_and_wait_for_image, | |
| inputs=[api_key, input_image, garment_type], | |
| outputs=[output_image, status_output] | |
| ) | |
| cancel_btn.click( | |
| fn=cancel_request, | |
| inputs=None, | |
| outputs=None, | |
| queue=False | |
| ) | |
| gr.Markdown(""" | |
| # Note:\n | |
| The image will be resized to 1280px on the short edge. | |
| """) | |
| if __name__ == "__main__": | |
| # Verify required files exist before launching | |
| try: | |
| validate_assets() | |
| demo.launch() | |
| except FileNotFoundError as e: | |
| print(f"Error: {str(e)}") | |
| print("Please make sure these files exist in the same directory:") | |
| print("- logo.png") | |
| print("- pullover.png (for top selection)") | |
| print("- sweatpants.png (for bottom selection)") |