Upload 3 files

app.py

@@ -0,0 +1,57 @@
+import os
+import subprocess
+import gradio as gr
+from caption import generate_caption  # Your BLIP-2 captioning script
+
+# ===== 1. Install Kohya_SS Manually =====
+if not os.path.exists("kohya_ss"):
+    print("⬇️ Cloning Kohya_SS...")
+    os.system("git clone https://github.com/bmaltais/kohya_ss")
+    os.chdir("kohya_ss")
+    os.system("pip install -r requirements.txt")
+    os.system("pip install .")  # Editable install
+    os.chdir("..")
+
+# ===== 2. Training Function =====
+def train_lora(images, trigger_word, model_choice="Flux"):
+    # Save images
+    os.makedirs("train", exist_ok=True)
+    for i, img in enumerate(images):
+        img_path = f"train/img_{i}.jpg"
+        img.save(img_path)
+        # Auto-caption (from caption.py)
+        caption = generate_caption(img_path, trigger_word)
+        with open(f"train/img_{i}.txt", "w") as f:
+            f.write(caption)
+
+    # Train LoRA (simplified Kohya_SS command)
+    cmd = f"""
+    python kohya_ss/train_network.py \
+        --pretrained_model_name_or_path="{model_choice}" \
+        --train_data_dir="train" \
+        --output_dir="output" \
+        --resolution=512 \
+        --network_dim=64 \
+        --lr=1e-4 \
+        --max_train_steps=1000
+    """
+    subprocess.run(cmd, shell=True)
+    return "output/lora.safetensors"
+
+# ===== 3. Gradio UI =====
+with gr.Blocks() as demo:
+    gr.Markdown("## 🎨 1-Click LoRA Trainer")
+    with gr.Row():
+        images = gr.Files(label="Upload 30 Images", file_types=["image"])
+        trigger = gr.Textbox(label="Trigger Word (e.g., 'my_char')")
+    train_btn = gr.Button("🚀 Train LoRA")
+    output = gr.File(label="Download LoRA")
+    
+    train_btn.click(
+        train_lora,
+        inputs=[images, trigger],
+        outputs=output
+    )
+
+if __name__ == "__main__":
+    demo.launch()

caption.py

@@ -0,0 +1,21 @@
+from transformers import Blip2Processor, Blip2ForConditionalGeneration
+from PIL import Image
+import torch
+
+def generate_caption(image_path, trigger_word):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # Load BLIP-2 (smaller model for HF Spaces)
+    processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
+    model = Blip2ForConditionalGeneration.from_pretrained(
+        "Salesforce/blip2-opt-2.7b", 
+        torch_dtype=torch.float16
+    ).to(device)
+    
+    # Generate caption
+    image = Image.open(image_path)
+    inputs = processor(image, return_tensors="pt").to(device, torch.float16)
+    generated_ids = model.generate(**inputs, max_new_tokens=50)
+    caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+    
+    return f"a photo of [{trigger_word}], {caption}"

requirements.txt

@@ -0,0 +1,7 @@
+gradio==3.50.2
+torch==2.2.1
+accelerate==0.27.2
+diffusers==0.27.2
+transformers==4.38.2
+huggingface-hub==0.22.2
+xformers==0.0.23.post1