Create app.py

app.py

@@ -0,0 +1,268 @@
+import gradio as gr
+import os
+import shutil
+import tempfile
+import zipfile
+from pathlib import Path
+from datasets import load_dataset
+from huggingface_hub import HfApi, login
+from PIL import Image
+
+
+def process_hf_dataset(dataset_name: str, image_col: str, caption_col: str, lora_name: str, progress=gr.Progress()):
+    """Process a HuggingFace dataset and create image/txt pairs."""
+    if not dataset_name.strip():
+        return None, "Please enter a dataset name", lora_name
+    
+    try:
+        progress(0, desc="Loading dataset...")
+        ds = load_dataset(dataset_name, split="train")
+        
+        # Create temp directory for output
+        output_dir = tempfile.mkdtemp()
+        
+        # Detect columns if not specified
+        if not image_col:
+            for col in ds.column_names:
+                if ds.features[col].dtype == "image" or "image" in col.lower():
+                    image_col = col
+                    break
+        
+        if not caption_col:
+            for col in ds.column_names:
+                if "text" in col.lower() or "caption" in col.lower() or "prompt" in col.lower():
+                    caption_col = col
+                    break
+        
+        if not image_col or not caption_col:
+            return None, f"Could not detect columns. Available: {ds.column_names}", lora_name
+        
+        progress(0.1, desc=f"Processing {len(ds)} images...")
+        
+        for i, item in enumerate(ds):
+            progress((i + 1) / len(ds), desc=f"Processing image {i+1}/{len(ds)}")
+            
+            # Save image
+            img = item[image_col]
+            if not isinstance(img, Image.Image):
+                img = Image.open(img)
+            
+            img_filename = f"{i:05d}.png"
+            txt_filename = f"{i:05d}.txt"
+            
+            img.save(os.path.join(output_dir, img_filename))
+            
+            # Save caption
+            caption = item[caption_col]
+            with open(os.path.join(output_dir, txt_filename), "w", encoding="utf-8") as f:
+                f.write(str(caption))
+        
+        return output_dir, f"Processed {len(ds)} images from {dataset_name}", lora_name
+    
+    except Exception as e:
+        return None, f"Error: {str(e)}", lora_name
+
+
+def process_uploaded_images(images: list, caption: str, lora_name: str, progress=gr.Progress()):
+    """Process uploaded images with a shared caption."""
+    if not images:
+        return None, "Please upload some images", lora_name
+    
+    output_dir = tempfile.mkdtemp()
+    
+    for i, img_data in enumerate(progress.tqdm(images, desc="Processing images")):
+        # Gallery returns tuples of (filepath, caption) or just filepath
+        if isinstance(img_data, tuple):
+            img_path = img_data[0]
+        else:
+            img_path = img_data
+        
+        img = Image.open(img_path)
+        
+        # Use original filename without extension
+        orig_name = Path(img_path).stem
+        img_filename = f"{orig_name}.png"
+        txt_filename = f"{orig_name}.txt"
+        
+        img.save(os.path.join(output_dir, img_filename))
+        
+        with open(os.path.join(output_dir, txt_filename), "w", encoding="utf-8") as f:
+            f.write(caption if caption else "")
+    
+    return output_dir, f"Processed {len(images)} images", lora_name
+
+
+def create_zip(output_dir: str, lora_name: str = None):
+    """Create a zip file from the output directory."""
+    if not output_dir or not os.path.exists(output_dir):
+        return None
+    
+    # Use lora_name for zip filename if provided
+    if lora_name and lora_name.strip():
+        zip_filename = f"{lora_name.strip().replace(' ', '_')}.zip"
+        zip_path = os.path.join(tempfile.gettempdir(), zip_filename)
+    else:
+        zip_path = tempfile.mktemp(suffix=".zip")
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
+        for file in os.listdir(output_dir):
+            zf.write(os.path.join(output_dir, file), file)
+    
+    return zip_path
+
+
+def push_to_hub(output_dir: str, repo_name: str, token: str, private: bool, progress=gr.Progress()):
+    """Push the processed dataset to HuggingFace Hub."""
+    if not output_dir or not os.path.exists(output_dir):
+        return "No data to push. Process a dataset first."
+    
+    if not repo_name or not repo_name.strip():
+        return "Please enter a repository name (or provide a LoRA name when processing)"
+    
+    if not token or not token.strip():
+        return "Please enter your HuggingFace token"
+    
+    try:
+        progress(0, desc="Logging in...")
+        api = HfApi(token=token)
+        
+        progress(0.2, desc="Creating repository...")
+        api.create_repo(repo_name, repo_type="dataset", private=private, exist_ok=True)
+        
+        progress(0.4, desc="Uploading files...")
+        api.upload_folder(
+            folder_path=output_dir,
+            repo_id=repo_name,
+            repo_type="dataset",
+        )
+        
+        return f"Successfully pushed to https://huggingface.co/datasets/{repo_name}"
+    
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+
+# Global state for output directory
+current_output_dir = {"path": None, "lora_name": None}
+
+
+def process_dataset_wrapper(dataset_name, image_col, caption_col, lora_name, progress=gr.Progress()):
+    output_dir, msg, lora = process_hf_dataset(dataset_name, image_col, caption_col, lora_name, progress)
+    current_output_dir["path"] = output_dir
+    current_output_dir["lora_name"] = lora
+    zip_path = create_zip(output_dir, lora) if output_dir else None
+    return msg, zip_path
+
+
+def process_images_wrapper(images, caption, lora_name, progress=gr.Progress()):
+    output_dir, msg, lora = process_uploaded_images(images, caption, lora_name, progress)
+    current_output_dir["path"] = output_dir
+    current_output_dir["lora_name"] = lora
+    zip_path = create_zip(output_dir, lora) if output_dir else None
+    return msg, zip_path
+
+
+def push_wrapper(repo_name, token, private, progress=gr.Progress()):
+    # Use lora_name as default repo name if repo_name is empty
+    final_repo_name = repo_name.strip() if repo_name.strip() else None
+    if not final_repo_name and current_output_dir["lora_name"]:
+        final_repo_name = current_output_dir["lora_name"].strip().replace(" ", "_")
+    return push_to_hub(current_output_dir["path"], final_repo_name, token, private, progress)
+
+
+# Build the Gradio interface
+with gr.Blocks(title="AI Toolkit Dataset Converter") as demo:
+    gr.Markdown("# AI Toolkit Dataset Converter")
+    gr.Markdown("""Convert your datasets to the format expected by [ostris AI Toolkit](https://github.com/ostris/ai-toolkit?tab=readme-ov-file#dataset-preparation). You can either: 
+    1. provide a dataset name from the hub OR
+    2. upload your images directly
+""")
+    
+    with gr.Tabs():
+        # Tab 1: HuggingFace Dataset
+        with gr.Tab("From HuggingFace Dataset"):
+            dataset_name = gr.Textbox(
+                label="Dataset Name",
+                placeholder="e.g., Norod78/Yarn-art-style"
+            )
+            lora_name_ds = gr.Textbox(
+                label="LoRA Name (optional)",
+                placeholder="e.g., my-lora-style",
+                info="Used for ZIP filename and Hub dataset name"
+            )
+            with gr.Row():
+                image_col = gr.Textbox(
+                    label="Image Column (leave empty to auto-detect)",
+                    placeholder="image"
+                )
+                caption_col = gr.Textbox(
+                    label="Caption Column (leave empty to auto-detect)",
+                    placeholder="text"
+                )
+            process_ds_btn = gr.Button("Process Dataset", variant="primary")
+            ds_status = gr.Textbox(label="Status", interactive=False)
+            ds_download = gr.File(label="Download ZIP")
+            
+            process_ds_btn.click(
+                process_dataset_wrapper,
+                inputs=[dataset_name, image_col, caption_col, lora_name_ds],
+                outputs=[ds_status, ds_download]
+            )
+        
+        # Tab 2: Upload Images
+        with gr.Tab("From Uploaded Images"):
+            images_input = gr.Gallery(
+                label="Upload Images",
+                file_types=["image"],
+                interactive=True,
+                columns=4,
+                height="auto"
+            )
+            lora_name_img = gr.Textbox(
+                label="LoRA Name (optional)",
+                placeholder="e.g., my-lora-style",
+                info="Used for ZIP filename and Hub dataset name"
+            )
+            shared_caption = gr.Textbox(
+                label="Caption for all images",
+                placeholder="Enter a caption to use for all uploaded images",
+                lines=3
+            )
+            process_img_btn = gr.Button("Process Images", variant="primary")
+            img_status = gr.Textbox(label="Status", interactive=False)
+            img_download = gr.File(label="Download ZIP")
+            
+            process_img_btn.click(
+                process_images_wrapper,
+                inputs=[images_input, shared_caption, lora_name_img],
+                outputs=[img_status, img_download]
+            )
+    
+    # Push to Hub section
+    gr.Markdown("---")
+    gr.Markdown("### Push to HuggingFace Hub")
+    
+    with gr.Row():
+        repo_name = gr.Textbox(
+            label="Repository Name",
+            placeholder="username/dataset-name (uses LoRA name if empty)",
+            info="Leave empty to use LoRA name as dataset name"
+        )
+        hf_token = gr.Textbox(
+            label="HuggingFace Token",
+            type="password",
+            placeholder="hf_..."
+        )
+    
+    private_repo = gr.Checkbox(label="Private Repository", value=False)
+    push_btn = gr.Button("Push to Hub", variant="secondary")
+    push_status = gr.Textbox(label="Push Status", interactive=False)
+    
+    push_btn.click(
+        push_wrapper,
+        inputs=[repo_name, hf_token, private_repo],
+        outputs=[push_status]
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()