Update app.py

app.py

@@ -1,9 +1,8 @@
-import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline
 import gradio as gr
+import spaces
 import random
-import re
 
 model_id = "kkntr/sdxl-kkntr"
 
@@ -13,8 +12,18 @@ pipe = StableDiffusionXLPipeline.from_pretrained(
     use_safetensors=True
 ).to("cuda")
 
-# Scratchpad-listat
-scratchpads = {
+pipe.load_lora_weights(
+    {
+        "kkntr/lora1": 0.2,
+        "kkntr/lora2": 0.3,
+        "kkntr/lora12": 0.4,
+        "kkntr/lora10": 0.1,
+    }
+)
+pipe.fuse_lora()
+
+# Scratchpad vaihtoehdot
+scratchpad = {
     "animal": [
         "kangaroo", "otter", "tiger", "lion", "wolf", "fox", "polar bear", "deer", "reindeer", "hyena", "bull", "rabbit", "panther", "moose", "cheetah", "badger", "striped hyena", "rat", "crash bandicoot", "leopard", "thylacine", "bighorn sheep", "zebra", "horse", "donkey", "unicorn, white fur, rainbow hair, mythological equine", "alligator", "western dragon", "lizard", "crocodile", "theropod", "montgomery gator, fnaf", "tagg, dinosaur, reptile, scalie, theropod", "black russian terrier", "pit bull", "schnauzer", "dobermann", "mastiff", "molosser", "pinscher", "rottweiler", "tamaskan dog", "wolfdog", "alaskan husky", "alaskan malamute", "canadian eskimo dog", "german shepherd", "labrador husky", "siberian husky", "labrador", "golden retriever", "dalmatian", "border collie", "collie", "herding dog", "pastoral dog", "sheepdog", "carolina dog", "rhodesian ridgeback", "nordic sled dog", "spitz", "husky", "akita", "mountain dog", "primitive dog", "siberian retriever", "hunting dog", "australian cattle dog",
     ],
@@ -38,22 +47,60 @@ scratchpads = {
     ],
 }
 
-def expand_scratchpad_tags(prompt: str) -> str:
-    """
-    Korvaa [tagi] scratchpad-listasta satunnaisella vaihtoehdolla.
-    """
-    def replacer(match):
-        tag = match.group(1)
-        if tag in scratchpads:
-            return random.choice(scratchpads[tag])
-        return match.group(0)  # jätä koskematta jos ei löydy
+def expand_prompt(prompt: str) -> str:
+    """Korvaa scratchpad-tagit satunnaisilla vaihtoehdoilla"""
+    for tag, options in scratchpad.items():
+        if tag in prompt:
+            choice = random.choice(options)
+            prompt = prompt.replace(tag, choice, 1)  # korvataan vain kerran
+    return prompt
 
-    return re.sub(r"\[([a-zA-Z0-9_]+)\]", replacer, prompt)
+resolutions = [
+    (592, 1768), (608, 1696), (640, 1624), (672, 1544),
+    (720, 1440), (744, 1392), (784, 1320), (832, 1248),
+    (864, 1208), (880, 1184), (912, 1144), (976, 1072),
+    (1024, 1024), (1072, 912), (1144, 912), (1184, 880),
+    (1208, 864), (1248, 832), (1320, 784), (1392, 744),
+    (1440, 720), (1544, 672), (1624, 640), (1696, 608),
+    (1768, 592),
+]
+
+lora_styles = {
+    "None": (None, 0.0, ""),
+    "Foreskin play": ("kkntr/lora5", 1.2, "foreskin insertion"),
+    "Cum through clothing": ("kkntr/lora6", 1.2, "cum in clothing"),
+    "Penis through leghole": ("kkntr/lora7", 9.0, "(((penis through leghole))), (loose shorts), (penis outline), poking out, genital outline, bulge flowing sideways along thigh, visible penis shape under fabric, tight crotch, fabric tension"),
+    "Frottage": ("kkntr/lora8", 1.2, "byfrottage, frottage"),
+    "Penis towards viewer": ("kkntr/lora9", 1.2, "penis_towards_viewer"),
+    "Foreskin play closeup": ("kkntr/lora11", 2.0, "ZonkFS"),
+}
 
 @spaces.GPU
-def generate(prompt, negative_prompt, width, height):
-    # Korvataan scratchpad-tagit ennen kuin prompt menee pipeen
-    prompt = expand_scratchpad_tags(prompt)
+def generate(prompt, negative_prompt, width, height, random_res, selected_style):
+    if random_res:
+        width, height = random.choice(resolutions)
+
+    # Expandataan scratchpad-tagit
+    prompt = expand_prompt(prompt)
+
+    # Valitaan tyyli-LoRA ja lisätään triggeri
+    lora_addition = {}
+    if selected_style != "None" and selected_style in lora_styles:
+        lora_path, lora_strength, lora_trigger = lora_styles[selected_style]
+        prompt = f"{lora_trigger}, {prompt}"
+        lora_addition[lora_path] = lora_strength
+
+    # Ladataan pysyvät + tyyli-LoRA yhteen
+    all_loras = {
+        "kkntr/lora1": 0.2,
+        "kkntr/lora2": 0.3,
+        "kkntr/lora12": 0.4,
+        "kkntr/lora10": 0.1,
+    }
+    all_loras.update(lora_addition)
+
+    pipe.load_lora_weights(all_loras)
+    pipe.fuse_lora()
 
     quality_suffix = (
         ", (anthro), (male), (best quality), high resolution, (ultra detailed),"
@@ -77,30 +124,27 @@ def generate(prompt, negative_prompt, width, height):
 
     return image, full_prompt
 
-# Gradio UI
 gr.Interface(
     fn=generate,
     inputs=[
-        gr.Textbox(label="Prompt"),
+        gr.Textbox(label="Prompt", value="a wolf [kuvakulma], [tunnelma], [värit]"),
         gr.Textbox(
             label="Negative Prompt",
-            value="(worst quality), (low quality), (blurry), (bad anatomy), "
-                  "(fused fingers), (extra fingers), (missing fingers), "
-                  "(deformed), (jpeg artifacts), (watermark), (signature), "
-                  "text, error, out of frame, bad hands, bad feet, bad proportions, "
-                  "fused limbs, cropped, username, logo, bad perspective, "
-                  "lowres, overexposure, bad artist, bad composition, "
-                  "distorted, oversaturated, mutated, bad shadow, multiple limbs, "
-                  "extra arms, extra legs, bad background, tilted, bad lighting, "
-                  "bad face, ugly, bad eyes, disfigured, multiple pictures"
+            value="(worst quality), (low quality), (blurry), (bad anatomy), (fused fingers), (extra fingers), (missing fingers), (deformed), (jpeg artifacts), (watermark), (signature), text, error, out of frame, bad hands, bad feet, bad proportions, fused limbs, cropped, username, logo, bad perspective, lowres, overexposure, bad artist, bad composition, distorted, oversaturated, mutated, bad shadow, multiple limbs, extra arms, extra legs, bad background, tilted, bad lighting, bad face, ugly, bad eyes, disfigured, multiple pictures"
         ),
         gr.Slider(minimum=512, maximum=1024, step=64, value=768, label="Width"),
         gr.Slider(minimum=512, maximum=1024, step=64, value=768, label="Height"),
+        gr.Checkbox(label="Random resolution mode (override sliders)", value=True),
+        gr.Dropdown(
+            choices=list(lora_styles.keys()),
+            value="None",
+            label="Select Style"
+        )
     ],
     outputs=[
         gr.Image(label="Generated Image"),
         gr.Textbox(label="Used Prompt", lines=4)
     ],
-    title="FURRYSTYLE Scratchpad",
-    description="animal, pose, swimwear, underwear, topwear, bottomwear, colour"
+    title="FURRYSTYLE",
+    description="Give a prompt (you can use [kuvakulma], [tunnelma], [värit] for random scratchpad expansion)"
 ).launch()