Upload 2 files

app.py

@@ -18,18 +18,61 @@ pipe2 = AutoPipelineForText2Image.from_pretrained(
 
 @spaces.GPU
 def infer(
-        prompt,
-        negative_prompt,
-        seed,
-        randomize_seed,
-        width,
-        height,
-        guidance_scale,
-        num_inference_steps,
-        model_name,
-        checkpoint,
+        prompt: str,
+        negative_prompt: str,
+        seed: int,
+        randomize_seed: bool,
+        width: int,
+        height: int,
+        guidance_scale: float,
+        num_inference_steps: int,
+        model_name: str,
+        checkpoint: str,
         progress=gr.Progress(track_tqdm=True),
 ):
+    """
+    Generate an image using the specified parameters and LoRA checkpoint.
+    
+    This function loads a LoRA checkpoint and generates an image based on the provided
+    prompt and generation parameters. It supports both Stable Diffusion v1.5 and SDXL models.
+    
+    Parameters
+    ----------
+    prompt : str
+        Text prompt describing the image to generate
+    negative_prompt : str
+        Text prompt describing what to avoid in the generated image
+    seed : int
+        Random seed for reproducible generation
+    randomize_seed : bool
+        Whether to use a random seed instead of the provided one
+    width : int
+        Width of the generated image
+    height : int
+        Height of the generated image
+    guidance_scale : float
+        Strength of prompt guidance (higher values = stronger adherence to prompt)
+    num_inference_steps : int
+        Number of denoising steps for generation
+    model_name : str
+        Name of the base model to use
+    checkpoint : str
+        Name of the LoRA checkpoint to load
+    progress : gr.Progress, optional
+        Gradio progress tracker
+        
+    Returns
+    -------
+    tuple
+        Tuple containing (generated_image, seed_used)
+        
+    Notes
+    -----
+    - Automatically selects the appropriate pipeline based on model_name
+    - Loads LoRA weights from the checkpoints directory
+    - Uses torch.Generator for reproducible results
+    - Returns both the generated image and the seed used
+    """
     pipe = pipe2 if model_name == 'runwayml/stable-diffusion-v1-5' else pipe1
     pipe.load_lora_weights(os.path.join('checkpoints', checkpoint),
                            weight_name="pytorch_lora_weights.safetensors")
@@ -52,10 +95,17 @@ def infer(
 
 
 def main():
+    """
+    Create and configure the Gradio web interface.
+    
+    This function sets up the complete Gradio interface with all necessary components
+    for image generation, including input controls, output display, and example prompts.
+
+    """
     examples = [
-        "Ioannis hiking in the mountains",
+        "Ioannis hiking in the mountains looking at the camera",
         "Ioannis at an office desk",
-        "Ioannis skiing in the Alps",
+        "Ioannis skiing in the Alps looking at the camera without mask",
     ]
 
     css = """
@@ -140,7 +190,7 @@ def main():
                         minimum=1,
                         maximum=400,
                         step=1,
-                        value=150,  # Replace with defaults that work for your checkpoints
+                        value=200,  # Replace with defaults that work for your checkpoints
                     )
 
                 with gr.Row():

config.py

@@ -2,17 +2,22 @@ import torch
 import numpy as np
 import os
 
-
+# Device configuration
 device = "cuda" if torch.cuda.is_available() else "mps"
 torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
+# Constants
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-PROJECT_PATH = "/Users/ioannisxarchakos/Dropbox/Projects/Dyla/"
+# Project paths
+PROJECT_PATH = ""
+
+# Model configurations
 PRETRAINED_MODEL = 'runwayml/stable-diffusion-v1-5'
 PRETRAINED_MODEL_XL = 'stabilityai/stable-diffusion-xl-base-1.0'
 LORA_WEIGHTS_PATH = os.path.join(PROJECT_PATH, 'sd-checkpoints-finetuned-lora/')
-DATA_DIR = os.path.join(PROJECT_PATH, 'ioannis_images/')
+IMAGE_FOLDER = "ioannis_images"
+DATA_DIR = os.path.join(PROJECT_PATH, IMAGE_FOLDER)
 CSV_PATH = os.path.join(DATA_DIR, 'prompts.csv')
 OUTPUT_DIR = 'generated_images'