Update app.py

app.py

@@ -5,12 +5,14 @@ import torch
 import os
 import json
 from datetime import datetime
-from PIL import Image
 
+# Path to save/load preset prompts
 PRESET_PATH = "prompt_presets.json"
 IMAGE_SAVE_DIR = "generated_images"
 os.makedirs(IMAGE_SAVE_DIR, exist_ok=True)
 
+# Optional: Upload to HF Hub with repo check
+
 def upload_to_huggingface(image_path, token, repo_id):
     try:
         from huggingface_hub import HfApi
@@ -32,7 +34,7 @@ def upload_to_huggingface(image_path, token, repo_id):
     except Exception as e:
         return f"❌ Upload failed: {e}"
 
-def start_training(instance_token, class_token, zip_file, output_dir, max_steps, lr, hf_token):
+def start_training(instance_token, class_token, zip_file, output_dir, max_steps, lr, hf_token, seed, precision):
     return train_model(
         instance_token=instance_token,
         class_token=class_token,
@@ -40,10 +42,12 @@ def start_training(instance_token, class_token, zip_file, output_dir, max_steps,
         output_dir=output_dir,
         max_train_steps=max_steps,
         learning_rate=lr,
-        hf_token=hf_token
+        hf_token=hf_token,
+        seed=seed,
+        precision=precision
     )
 
-def generate_image(prompt, hf_token, model_dir, hf_dataset_repo, reference_img=None):
+def generate_image(prompt, hf_token, model_dir, hf_dataset_repo):
     try:
         if not os.path.exists(model_dir):
             return "❌ Error: Trained model directory not found. Please train your model first.", None, None, None
@@ -53,10 +57,6 @@ def generate_image(prompt, hf_token, model_dir, hf_dataset_repo, reference_img=N
             torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
         ).to("cuda" if torch.cuda.is_available() else "cpu")
 
-        if reference_img:
-            ref = Image.open(reference_img).convert("RGB")
-            prompt += " using the pose/style from reference image"
-
         image = pipe(prompt).images[0]
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         save_path = os.path.join(IMAGE_SAVE_DIR, f"generated_{timestamp}.png")
@@ -66,76 +66,24 @@ def generate_image(prompt, hf_token, model_dir, hf_dataset_repo, reference_img=N
     except Exception as e:
         return f"❌ Generation failed: {str(e)}", None, None, None
 
-def build_character_prompt(**kwargs):
-    prompt_parts = []
-    for key, val in kwargs.items():
-        if val:
-            prompt_parts.append(f"{key.replace('_', ' ').title()}: {val}")
-    return "\n".join(prompt_parts)
-
 def create_ui():
     with gr.Blocks() as demo:
-        with gr.Tab("Character Builder"):
-            with gr.Row():
-                species = gr.Dropdown(label="Species", choices=["human", "animal hybrid", "demon", "spirit", "elf", "dwarf", "vampire", "fairy", "elemental", "monster"], value="human")
-                gender = gr.Dropdown(label="Gender", choices=["male", "female", "trans male", "trans female", "random"], value="female")
-                skin_tone = gr.Dropdown(label="Skin Tone", choices=["white", "grey", "black", "dark brown", "light brown", "deep tan", "golden tan", "light tan", "beige", "pale", "random"], value="beige")
-            with gr.Row():
-                head_shape = gr.Dropdown(label="Head Shape", choices=["oval", "round", "square", "diamond", "heart", "triangle", "random"])
-                right_eye = gr.Dropdown(label="Right Eye Color", choices=["black", "brown", "hazel", "green", "blue", "yellow", "red", "amber", "orange", "purple", "random"])
-                left_eye = gr.Dropdown(label="Left Eye Color", choices=["black", "brown", "hazel", "green", "blue", "yellow", "red", "amber", "orange", "purple", "random"])
-            with gr.Row():
-                eye_shape = gr.Dropdown(label="Eye Shape", choices=["almond", "round", "protruding", "mono-lid", "downturn", "upturn", "close-set", "wide-set", "deep-set", "hooded", "random"])
-                nose_shape = gr.Dropdown(label="Nose Shape", choices=["hooked", "droopy", "Roman", "aquiline", "Grecian", "button", "pixie", "upturned", "snub", "funnel", "flat", "wide", "bulbous", "crooked", "random"])
-                lip_shape = gr.Dropdown(label="Lip Shape", choices=["full", "wide", "narrow", "full upper", "full lower", "heart", "bow", "round", "thin", "downturned", "random"])
-            with gr.Row():
-                height = gr.Dropdown(label="Height", choices=["very short", "short", "average", "tall", "very tall", "gigantic", "enormous", "random"])
-                body_type = gr.Dropdown(label="Body Type Female", choices=["triangle", "inverted triangle", "rectangle", "oval", "round", "hourglass", "bottom heavy", "top heavy", "apple shaped", "pear shaped", "diamond", "thin", "lean", "athletic", "curvy", "voluptuous", "thick", "muscular", "petite", "fat", "random"], multiselect=True)
-
-            torso = gr.Textbox(label="Torso Clothing Description")
-            position = gr.Textbox(label="Body Position")
-            background = gr.Textbox(label="Background Description")
-            lighting = gr.Textbox(label="Lighting Description")
-            mood = gr.Textbox(label="Mood")
-
-            reference_img = gr.Image(label="Reference Image (Optional)", type="filepath")
-            randomize_btn = gr.Button("Randomize")
-            art_style = gr.Dropdown(label="Art Style Preset", choices=["Fantasy Manga", "Golden Age Comic", "Silver Age Comic", "90’s Extreme Comic Book", "Modern Comic Book", "Anime", "Realistic"], value="Fantasy Manga")
-            hf_token = gr.Textbox(label="Hugging Face Token", type="password")
-            model_dir = gr.Textbox(label="Trained Model Directory", value="trained_model")
-            hf_dataset_repo = gr.Textbox(label="Hugging Face Dataset Repo ID", value="generated-images")
-            prompt_output = gr.Textbox(label="Generated Prompt")
-            generate_btn = gr.Button("Generate Prompt & Image")
-            image_output = gr.Image(label="Generated Image")
-            file_path_display = gr.Textbox(label="File Path")
-            upload_status = gr.Textbox(label="Upload Status")
-
-            def compose_and_generate(*args):
-                keys = ["species", "gender", "skin_tone", "head_shape", "right_eye", "left_eye", "eye_shape", "nose_shape", "lip_shape", "height", "body_type", "torso", "position", "background", "lighting", "mood", "art_style"]
-                prompt = build_character_prompt(**{k: v for k, v in zip(keys, args[:-4]) if v})
-                prompt += f"\nArt style: {args[-4]}"
-                return (prompt,) + generate_image(prompt, args[-3], args[-2], args[-1], args[-5])
-
-            generate_btn.click(
-                fn=compose_and_generate,
-                inputs=[species, gender, skin_tone, head_shape, right_eye, left_eye, eye_shape, nose_shape, lip_shape, height, body_type, torso, position, background, lighting, mood, art_style, reference_img, hf_token, model_dir, hf_dataset_repo],
-                outputs=[prompt_output, image_output, file_path_display, upload_status]
-            )
-
         with gr.Tab("Train Model"):
             instance_token = gr.Textbox(label="Instance Token")
             class_token = gr.Textbox(label="Class Token")
             zip_file = gr.File(label="Training ZIP File")
             output_dir = gr.Textbox(label="Output Directory", value="trained_model")
-            max_steps = gr.Number(label="Max Training Steps", value=400)
+            max_steps = gr.Number(label="Max Training Steps", value=1200)
             lr = gr.Number(label="Learning Rate", value=5e-6)
+            seed = gr.Number(label="Random Seed", value=42)
+            precision = gr.Dropdown(label="Precision Mode", choices=["fp16", "fp32"], value="fp16")
             hf_token_train = gr.Textbox(label="Hugging Face Token", type="password")
             train_btn = gr.Button("Start Training")
-            train_output = gr.Textbox(label="Training Output")
+            train_output = gr.Textbox(label="Training Output", lines=8)
 
             train_btn.click(
                 fn=start_training,
-                inputs=[instance_token, class_token, zip_file, output_dir, max_steps, lr, hf_token_train],
+                inputs=[instance_token, class_token, zip_file, output_dir, max_steps, lr, hf_token_train, seed, precision],
                 outputs=train_output
             )
 
@@ -160,4 +108,3 @@ def create_ui():
 demo = create_ui()
 if __name__ == "__main__":
     demo.launch()
-