app1.py

import os
import requests
import zipfile
import gradio as gr
import shutil
from PIL import Image

# Configuration
API_URL = "https://porn-pictures-api.p.rapidapi.com/pornstars/{gender}/{page}"
API_HEADERS = {
    "x-rapidapi-key": "2dcd4142a4msh0d39ff26a5b144cp1e8703jsndf67b0582674",
    "x-rapidapi-host": "porn-pictures-api.p.rapidapi.com"
}
OUTPUT_DIR = "downloaded_images"  # Base output folder
IMAGES_DIR = os.path.join(OUTPUT_DIR, "images")  # Subfolder for downloaded images
ZIP_FILE = os.path.join(OUTPUT_DIR, "images.zip")  # Path for the output ZIP file

# Ensure output directory exists
os.makedirs(OUTPUT_DIR, exist_ok=True)

def fetch_image_urls(gender, num_images):
    """Fetch image URLs from the API based on gender and desired number of images."""
    image_urls = []
    page = 1
    items_per_page = 40  # From the response, each page has 40 items

    while len(image_urls) < num_images:
        url = API_URL.format(gender=gender.lower(), page=page)
        try:
            response = requests.get(url, headers=API_HEADERS)
            response.raise_for_status()  # Raise an error for bad status codes
            data = response.json()

            if "result" not in data or not data["result"]:
                break  # No more data to fetch

            for item in data["result"]:
                if len(image_urls) >= num_images:
                    break
                image_urls.append(item["picture"])
            
            page += 1
        except Exception as e:
            print(f"Error fetching page {page}: {e}")
            break
    
    return image_urls[:num_images]

def download_images(image_urls):
    """Download images from the provided URLs and save to IMAGES_DIR."""
    if os.path.exists(IMAGES_DIR):
        shutil.rmtree(IMAGES_DIR)  # Clear previous contents
    os.makedirs(IMAGES_DIR, exist_ok=True)

    downloaded_count = 0
    image_paths = []
    for idx, url in enumerate(image_urls, 1):
        try:
            response = requests.get(url, stream=True)
            response.raise_for_status()
            image_path = os.path.join(IMAGES_DIR, f"img{idx}.jpg")
            with open(image_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
            # Verify the image
            Image.open(image_path).verify()
            downloaded_count += 1
            image_paths.append(image_path)
            print(f"Downloaded {idx}/{len(image_urls)}: {url}")
        except Exception as e:
            print(f"Error downloading {url}: {e}")
    
    return downloaded_count, image_paths

def create_zip_file(selected_image_paths):
    """Create a ZIP file of the selected images."""
    if os.path.exists(ZIP_FILE):
        os.remove(ZIP_FILE)  # Remove previous ZIP
    with zipfile.ZipFile(ZIP_FILE, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for image_path in selected_image_paths:
            arcname = os.path.relpath(image_path, OUTPUT_DIR)
            zipf.write(image_path, arcname)
    return ZIP_FILE

def process_and_display(gender, num_images):
    """Fetch and download images, then prepare data for display."""
    num_images = int(num_images)  # Convert dropdown selection to integer
    if num_images > 24:
        num_images = 24  # Limit to 24 images to fit 6 rows of 4

    # Fetch image URLs
    image_urls = fetch_image_urls(gender, num_images)
    if not image_urls:
        return "Failed to fetch image URLs.", None, None, None

    # Download images
    downloaded_count, image_paths = download_images(image_urls)
    if downloaded_count == 0:
        return "No images were successfully downloaded.", None, None, None

    return f"Successfully downloaded {downloaded_count}/{num_images} images. Select images to include in ZIP below.", None, image_paths, image_paths

def process_zip_submission(image_paths, *checkbox_states):
    """Create a ZIP file based on the selected images."""
    if not image_paths:
        return "No images available to process.", None
    
    selected_image_paths = [image_paths[i] for i, state in enumerate(checkbox_states) if state]
    if not selected_image_paths:
        return "No images selected for ZIP.", None
    
    zip_path = create_zip_file(selected_image_paths)
    return f"ZIP file created with {len(selected_image_paths)} images at {zip_path}", zip_path

# Gradio Interface
with gr.Blocks(title="Image Downloader") as demo:
    gr.Markdown("### Select Parameters to Download Images")
    gender_input = gr.Dropdown(
        label="Category (Gender, Race, Hair Color, Body Type) - Select or Type Custom",
        choices=["female", "male", "nonbinary", "black", "latino", "asian", "white", "mixed",
                 "redhead", "blonde", "brunette", "blackhair", "bald", "bbw", "curvy",
                 "slim", "muscular", "thick"],
        value="female",
        allow_custom_value=True
    )
    num_images_input = gr.Dropdown(
        label="Number of Images (Max 24)",
        choices=["4", "8", "12", "16", "20", "24"],  # Multiples of 4 up to 24
        value="4"
    )
    download_button = gr.Button("Fetch and Display Images")
    
    gr.Markdown("### Download Status")
    status_output = gr.Textbox(label="Status", interactive=False)
    
    gr.Markdown("### Download Your Images")
    zip_output = gr.File(label="Download ZIP", visible=False)
    
    gr.Markdown("### Image Gallery (Click Thumbnails to View Full Size)")
    image_paths_state = gr.State()  # Store image paths for later use
    
    # Define constants for grid layout
    IMAGES_PER_ROW = 4
    MAX_ROWS = 6  # 6 rows of 4 = 24 images max
    TOTAL_IMAGES = IMAGES_PER_ROW * MAX_ROWS  # 24 images max
    
    # Create image and checkbox components in a strict 4-per-row grid
    image_outputs = []
    checkbox_outputs = []
    
    for row in range(MAX_ROWS):
        with gr.Row():
            for col in range(IMAGES_PER_ROW):
                idx = row * IMAGES_PER_ROW + col
                with gr.Column(min_width=150):  # Set a minimum width to force layout
                    image_output = gr.Image(
                        label=f"Image {idx+1}", 
                        visible=False, 
                        height=150, 
                        width=150
                    )
                    checkbox_output = gr.Checkbox(
                        label=f"Include in ZIP", 
                        value=True, 
                        visible=False
                    )
                    image_outputs.append(image_output)
                    checkbox_outputs.append(checkbox_output)
    
    gr.Markdown("### Submit Selections")
    submit_button = gr.Button("Create ZIP of Selected Images")
    
    def on_download(gender, num_images):
        status, zip_path, image_paths, _ = process_and_display(gender, num_images)
        if image_paths:
            # Update image and checkbox components for the downloaded images
            updates = {
                status_output: status,
                zip_output: gr.File(value=None, visible=False),  # Hide ZIP initially
                image_paths_state: image_paths
            }
            for i in range(TOTAL_IMAGES):
                if i < len(image_paths):
                    updates[image_outputs[i]] = gr.Image(
                        value=image_paths[i], 
                        visible=True, 
                        label=f"Image {i+1}", 
                        width=150, 
                        height=150
                    )
                    updates[checkbox_outputs[i]] = gr.Checkbox(
                        value=True, 
                        visible=True, 
                        label=f"Include in ZIP"
                    )
                else:
                    updates[image_outputs[i]] = gr.Image(value=None, visible=False)
                    updates[checkbox_outputs[i]] = gr.Checkbox(value=False, visible=False)
            return updates
        return {
            status_output: status,
            zip_output: gr.File(visible=False),
            image_paths_state: None,
            **{image_outputs[i]: gr.Image(value=None, visible=False) for i in range(TOTAL_IMAGES)},
            **{checkbox_outputs[i]: gr.Checkbox(value=False, visible=False) for i in range(TOTAL_IMAGES)}
        }
    
    def on_submit(image_paths, *checkbox_states):
        status, zip_path = process_zip_submission(image_paths, *checkbox_states)
        return {
            status_output: status,
            zip_output: gr.File(value=zip_path, visible=True) if zip_path else gr.File(visible=False)
        }
    
    download_button.click(
        fn=on_download,
        inputs=[gender_input, num_images_input],
        outputs=[status_output, zip_output, image_paths_state] + image_outputs + checkbox_outputs
    )
    
    submit_button.click(
        fn=on_submit,
        inputs=[image_paths_state] + checkbox_outputs,
        outputs=[status_output, zip_output]
    )

demo.launch()