Create app.py

app.py

@@ -0,0 +1,91 @@
+import os
+import shutil
+import subprocess
+import streamlit as st
+
+# ─── 1. Mode detection & data directory ───────────────────────────────────────
+# LOCAL_TRAIN=1 → use "./data"; otherwise Spaces uses "/tmp/data"
+LOCAL = os.environ.get("LOCAL_TRAIN", "").lower() in ("1", "true")
+DATA_DIR = os.path.join(os.getcwd(), "data") if LOCAL else "/tmp/data"
+os.makedirs(DATA_DIR, exist_ok=True)
+
+# ─── 2. Page layout ───────────────────────────────────────────────────────────
+st.set_page_config(page_title="HiDream LoRA Trainer", layout="wide")
+st.title("🎨 HiDream LoRA Trainer (Streamlit)")
+
+# Sidebar for configuration
+with st.sidebar:
+    st.header("🛠 Configuration")
+    base_model   = st.selectbox(
+        "Base Model",
+        ["HiDream-ai/HiDream-I1-Dev",
+         "runwayml/stable-diffusion-v1-5",
+         "stabilityai/stable-diffusion-2-1"]
+    )
+    trigger_word = st.text_input("Trigger Word", value="default-style")
+    num_steps    = st.slider("Training Steps", min_value=10, max_value=500, value=100, step=10)
+    lora_r       = st.slider("LoRA Rank (r)", min_value=4, max_value=128, value=16, step=4)
+    lora_alpha   = st.slider("LoRA Alpha", min_value=4, max_value=128, value=16, step=4)
+
+    st.markdown("---")
+    st.header("📂 Upload Dataset")
+    uploaded_files = st.file_uploader(
+        "Select your images & text files",
+        type=["jpg","jpeg","png","txt"],
+        accept_multiple_files=True
+    )
+    if st.button("Upload Dataset"):
+        # Clear old files
+        for f in os.listdir(DATA_DIR):
+            os.remove(os.path.join(DATA_DIR, f))
+        # Write new files
+        for up in uploaded_files:
+            dest = os.path.join(DATA_DIR, up.name)
+            with open(dest, "wb") as f:
+                f.write(up.getbuffer())
+        st.success(f"✅ Uploaded {len(uploaded_files)} files to `{DATA_DIR}`")
+
+    st.markdown("---")
+    # Trigger training
+    if st.button("🚀 Start Training"):
+        st.session_state.training = True
+
+# ─── 3. Training log area ─────────────────────────────────────────────────────
+log_area = st.empty()
+
+# ─── 4. Invoke training when triggered ────────────────────────────────────────
+if st.session_state.get("training", False):
+    st.info("Training started… Logs below:")
+    log_lines = []
+    # Prepare environment for train.py
+    env = os.environ.copy()
+    env.update({
+        "BASE_MODEL":    base_model,
+        "TRIGGER_WORD":  trigger_word,
+        "NUM_STEPS":     str(num_steps),
+        "LORA_R":        str(lora_r),
+        "LORA_ALPHA":    str(lora_alpha),
+        "LOCAL_TRAIN":   os.environ.get("LOCAL_TRAIN","")
+    })
+
+    # Launch train.py as subprocess and stream logs
+    proc = subprocess.Popen(
+        ["python3", "train.py"],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        env=env
+    )
+
+    for line in proc.stdout:
+        log_lines.append(line)
+        # Update the text area with all lines so far
+        log_area.text_area("Training Log", value="".join(log_lines), height=400)
+    proc.wait()
+
+    if proc.returncode == 0:
+        st.success("✅ Training complete!")
+    else:
+        st.error(f"❌ Training failed (exit code {proc.returncode})")
+    # Reset trigger
+    st.session_state.training = False