training-script

app.py

@@ -1,8 +1,11 @@
 import gradio as gr
 import numpy as np
 import random
-from diffusers import DiffusionPipeline
 import torch
+from diffusers import DiffusionPipeline, StableDiffusionXLBaseModel, StableDiffusionTrainer
+from transformers import CLIPTextModel, CLIPTokenizer, TrainingArguments
+from datasets import load_dataset
+from huggingface_hub import HfApi, HfFolder, Repository
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -11,7 +14,7 @@ if torch.cuda.is_available():
     pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
     pipe.enable_xformers_memory_efficient_attention()
     pipe = pipe.to(device)
-else: 
+else:
     pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
     pipe = pipe.to(device)
 
@@ -19,24 +22,77 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
-
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
         
     generator = torch.Generator().manual_seed(seed)
     
     image = pipe(
-        prompt = prompt, 
-        negative_prompt = negative_prompt,
-        guidance_scale = guidance_scale, 
-        num_inference_steps = num_inference_steps, 
-        width = width, 
-        height = height,
-        generator = generator
-    ).images[0] 
+        prompt=prompt, 
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale, 
+        num_inference_steps=num_inference_steps, 
+        width=width, 
+        height=height,
+        generator=generator
+    ).images[0]
     
     return image
 
+def get_latest_version(repo_id):
+    api = HfApi()
+    repo_info = api.repo_info(repo_id)
+    versions = [tag.name for tag in repo_info.tags]
+    if not versions:
+        return "v_0.0"
+    latest_version = sorted(versions)[-1]
+    return latest_version
+
+def increment_version(version):
+    major, minor = map(int, version.split('_')[1:])
+    minor += 1
+    return f"v_{major}.{minor}"
+
+def train_model(train_data_path, output_dir, num_train_epochs, per_device_train_batch_size, learning_rate):
+    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
+    text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+
+    base_model = StableDiffusionXLBaseModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
+
+    dataset = load_dataset('imagefolder', data_dir=train_data_path)
+
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        num_train_epochs=num_train_epochs,
+        per_device_train_batch_size=per_device_train_batch_size,
+        learning_rate=learning_rate,
+        logging_dir="./logs",
+        logging_steps=10,
+    )
+
+    trainer = StableDiffusionTrainer(
+        model=base_model,
+        args=training_args,
+        train_dataset=dataset['train'],
+        tokenizer=tokenizer,
+    )
+
+    trainer.train()
+    base_model.save_pretrained(output_dir)
+
+    # Publish the model
+    repo_id = "ZennyKenny/stable-diffusion-xl-base-1.0_NatalieDiffusion"
+    latest_version = get_latest_version(repo_id)
+    new_version = increment_version(latest_version)
+
+    api = HfApi()
+    token = HfFolder.get_token()
+    repo = Repository(output_dir, clone_from=repo_id, token=token)
+    repo.git_tag(new_version)
+    repo.push_tag(new_version)
+    
+    return f"Training complete. Model saved to {output_dir} and published as version {new_version}."
+
 examples = [
     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
     "An astronaut riding a green horse",
@@ -133,14 +189,54 @@ with gr.Blocks(css=css) as demo:
                 )
         
         gr.Examples(
-            examples = examples,
-            inputs = [prompt]
+            examples=examples,
+            inputs=[prompt]
+        )
+    
+    # Add new section for training the model
+    with gr.Accordion("Training Settings", open=False):
+        train_data_path = gr.Text(
+            label="Training Data Path",
+            placeholder="Enter the path to your training data",
+        )
+        output_dir = gr.Text(
+            label="Output Directory",
+            placeholder="Enter the output directory for the trained model",
         )
+        num_train_epochs = gr.Slider(
+            label="Number of Training Epochs",
+            minimum=1,
+            maximum=10,
+            step=1,
+            value=3,
+        )
+        per_device_train_batch_size = gr.Slider(
+            label="Batch Size per Device",
+            minimum=1,
+            maximum=16,
+            step=1,
+            value=4,
+        )
+        learning_rate = gr.Slider(
+            label="Learning Rate",
+            minimum=1e-5,
+            maximum=1e-3,
+            step=1e-5,
+            value=5e-5,
+        )
+        train_button = gr.Button("Train Model")
+        train_result = gr.Text(label="Training Result", show_label=False)
+
+    train_button.click(
+        fn=train_model,
+        inputs=[train_data_path, output_dir, num_train_epochs, per_device_train_batch_size, learning_rate],
+        outputs=[train_result],
+    )
 
     run_button.click(
-        fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs = [result]
+        fn=infer,
+        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        outputs=[result]
     )
 
-demo.queue().launch()
+demo.queue().launch()