Create main.py

main.py

@@ -0,0 +1,117 @@
+import os
+import json
+import shutil
+from pathlib import Path
+import torch
+import gradio as gr
+from diffusers import StableDiffusionPipeline, DDIMScheduler
+from transformers import CLIPTextModel, CLIPTokenizer
+from PIL import Image
+from torch import autocast
+
+# Define necessary paths and variables
+MODEL_NAME = "runwayml/stable-diffusion-v1-5"
+OUTPUT_DIR = "/output_model"
+INSTANCE_PROMPT = "photo of {identifier} person"
+CLASS_PROMPT = "photo of a person"
+SEED = 1337
+RESOLUTION = 512
+TRAIN_BATCH_SIZE = 1
+LEARNING_RATE = 1e-6
+MAX_TRAIN_STEPS = 800
+GUIDANCE_SCALE = 8.0
+NUM_INFERENCE_STEPS = 50
+
+# Function to fine-tune the model
+def fine_tune_model(instance_data_dir, identifier):
+    # Set up paths
+    instance_prompt = INSTANCE_PROMPT.format(identifier=identifier)
+    concepts_list = [
+        {
+            "instance_prompt": instance_prompt,
+            "class_prompt": CLASS_PROMPT,
+            "instance_data_dir": instance_data_dir,
+            "class_data_dir": "/sample_data/person"  # Placeholder for regularization images
+        }
+    ]
+
+    # Save concepts_list.json
+    with open("concepts_list.json", "w") as f:
+        json.dump(concepts_list, f, indent=4)
+
+    # Run the training script
+    os.system(f"""
+    python3 train_dreambooth.py \
+        --pretrained_model_name_or_path={MODEL_NAME} \
+        --output_dir={OUTPUT_DIR} \
+        --revision="fp16" \
+        --with_prior_preservation --prior_loss_weight=1.0 \
+        --seed={SEED} \
+        --resolution={RESOLUTION} \
+        --train_batch_size={TRAIN_BATCH_SIZE} \
+        --train_text_encoder \
+        --mixed_precision="fp16" \
+        --use_8bit_adam \
+        --gradient_accumulation_steps=1 \
+        --learning_rate={LEARNING_RATE} \
+        --max_train_steps={MAX_TRAIN_STEPS} \
+        --save_sample_prompt="{instance_prompt}" \
+        --concepts_list="concepts_list.json"
+    """)
+
+# Function for inference
+def generate_images(prompt, negative_prompt, num_samples, model_path, height=512, width=512, num_inference_steps=50, guidance_scale=7.5):
+    pipe = StableDiffusionPipeline.from_pretrained(model_path, safety_checker=None, torch_dtype=torch.float16).to("cuda")
+    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+    pipe.enable_xformers_memory_efficient_attention()
+    g_cuda = torch.Generator(device='cuda').manual_seed(SEED)
+
+    with torch.autocast("cuda"), torch.inference_mode():
+        images = pipe(
+            prompt,
+            height=height,
+            width=width,
+            negative_prompt=negative_prompt,
+            num_images_per_prompt=num_samples,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            generator=g_cuda
+        ).images
+
+    return images
+
+# Gradio UI function
+def inference_ui(identifier, prompt, negative_prompt, num_samples, height, width, num_inference_steps, guidance_scale):
+    model_path = OUTPUT_DIR
+    prompt = INSTANCE_PROMPT.format(identifier=identifier) + ", " + prompt
+    images = generate_images(prompt, negative_prompt, num_samples, model_path, height, width, num_inference_steps, guidance_scale)
+    return images
+
+# Define Gradio interface
+def create_gradio_ui():
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                identifier = gr.Textbox(label="Identifier", placeholder="Enter a unique identifier")
+                image_upload = gr.File(label="Upload Images", file_count="multiple", type="file")
+                finetune_button = gr.Button(value="Fine-Tune Model")
+                finetune_output = gr.Textbox(label="Fine-Tuning Output")
+
+            with gr.Column():
+                prompt = gr.Textbox(label="Prompt", value="photo of {identifier} person in a marriage hall")
+                negative_prompt = gr.Textbox(label="Negative Prompt", value="")
+                num_samples = gr.Number(label="Number of Samples", value=4)
+                guidance_scale = gr.Number(label="Guidance Scale", value=8)
+                height = gr.Number(label="Height", value=512)
+                width = gr.Number(label="Width", value=512)
+                num_inference_steps = gr.Slider(label="Steps", value=50)
+                generate_button = gr.Button(value="Generate Images")
+                gallery = gr.Gallery()
+
+            finetune_button.click(finetune_model, inputs=[image_upload, identifier], outputs=finetune_output)
+            generate_button.click(inference_ui, inputs=[identifier, prompt, negative_prompt, num_samples, height, width, num_inference_steps, guidance_scale], outputs=gallery)
+
+    demo.launch()
+
+if __name__ == "__main__":
+    create_gradio_ui()