Update app.py

app.py

@@ -7,7 +7,8 @@ import random
 from diffusers import FluxPipeline
 from translatepy import Translator
 from huggingface_hub import hf_hub_download
-
+import requests
+import re
 
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 translator = Translator()
@@ -31,11 +32,41 @@ JS = """function () {
 
 if torch.cuda.is_available():
     pipe = FluxPipeline.from_pretrained(model, torch_dtype=torch.bfloat16)
-    pipe.load_lora_weights("Shakker-Labs/FLUX.1-dev-LoRA-blended-realistic-illustration", weight_name="FLUX-dev-lora-blended_realistic_illustration.safetensors")
-    pipe.fuse_lora(lora_scale=1.0)
-    pipe.to(device="cuda", dtype=torch.bfloat16)
 
-# Function 
+
+def scrape_lora_link(url):
+    try:
+        # Send a GET request to the URL
+        response = requests.get(url)
+        response.raise_for_status()  # Raise an exception for bad status codes
+
+        # Get the content of the page
+        content = response.text
+
+        # Use regular expression to find the link
+        pattern = r'href="(.*?lora.*?\.safetensors\?download=true)"'
+        match = re.search(pattern, content)
+
+        if match:
+           safetensors_url = match.group(1)
+           filename = safetensors_url.split('/')[-1].split('?')[0]  # Extract the filename from the URL
+           return filename
+        else:
+            return None
+
+    except requests.RequestException as e:
+        print(f"An error occurred while fetching the URL: {e}")
+        return None
+
+def enable_lora(lora_scale, lora_in, lora_add):
+    if load_in:
+        lora_add = load_in
+    url = f'https://huggingface.co/{lora_add}/tree/main'
+    lora_name = scrape_lora_link(url)    
+    pipe.load_lora_weights(lora_add, weight_name=lora_name)
+    pipe.fuse_lora(lora_scale=lora_scale)
+    pipe.to(device="cuda", dtype=torch.bfloat16)
+    
 @spaces.GPU()
 def generate_image(
     prompt,
@@ -43,7 +74,7 @@ def generate_image(
     height=1024,
     scales=5,
     steps=24,
-    seed: int =-1,
+    seed=-1,
     nums=1,
     progress=gr.Progress(track_tqdm=True)):
 
@@ -72,6 +103,22 @@ def generate_image(
     return image, seed
 
 
+def gen(
+    prompt,
+    width=768,
+    height=1024,
+    scales=5,
+    steps=24,
+    seed=-1,
+    nums=1,
+    lora_scale=1.0, 
+    lora_in="", 
+    lora_add=""
+):
+    enable_lora(lora_scale, lora_in, lora_add)
+    generate_image(prompt,width,height,scales,steps,seed,nums)
+
+
 examples_bak = [
         "a female character with long, flowing hair that appears to be made of ethereal, swirling patterns resembling the Northern Lights or Aurora Borealis. The background is dominated by deep blues and purples, creating a mysterious and dramatic atmosphere. The character's face is serene, with pale skin and striking features. She wears a dark-colored outfit with subtle patterns. The overall style of the artwork is reminiscent of fantasy or supernatural genres",
         "Digital art, portrait of an anthropomorphic roaring Tiger warrior with full armor, close up in the middle of a battle, behind him there is a banner with the text \"Open Source\".", 
@@ -90,8 +137,10 @@ examples = [
     "A cartoon style European woman wearing glasses is eating a table of seafood,including lobster,oysters,and other shellfish,in a well lit modern restaurant. The background of the restaurant is very blurry,and she is holding the utensils ready to eat. There is a glass of red wine and various dishes on the table. The illustrations contrast with the real food and environment,creating a unique mixed media effect and high angle perspective",
     "A cartoon style European man opens his hands and takes a selfie under the Sydney Opera House",
     "Against the backdrop of the Eiffel Tower, a cartoon style European woman wearing a delicate white floral dress stands there, with the iconic building of the tower clearly visible under the azure sky, capturing the romantic charm of Paris. When she takes photos against this stunning background, her flowing hair adds a dreamy atmosphere",
-    "A cartoon-style Indian girl,dressed in traditional Indian costumes and veil,takes a selfie in front of the Taj Mahal,upper_body",
-    "A modern outdoor poolside scene,with a blonde European and American woman lying on an orange willow lounge chair. She is depicted as a cartoon character wearing a bikini and white sandals. The background features beige concrete walls,glass doors with curtains,and tropical palm trees. The sunlight casts shadows on the tile floor and building facade,creating a fresh and warm weather and a relaxed and peaceful atmosphere"
+    "close up portrait, Amidst the interplay of light and shadows in a photography studio,a soft spotlight traces the contours of a face,highlighting a figure clad in a sleek black turtleneck. The garment,hugging the skin with subtle luxury,complements the Caucasian model's understated makeup,embodying minimalist elegance. Behind,a pale gray backdrop extends,its fine texture shimmering subtly in the dim light,artfully balancing the composition and focusing attention on the subject. In a palette of black,gray,and skin tones,simplicity intertwines with profundity,as every detail whispers untold stories.",
+    "Caucasian,The image features a young woman of European descent standing in an studio setting,surrounded by silk. (She is wearing a silk dress),paired with a bold. Her brown hair is wet and tousled,falling naturally around her face,giving her a raw and edgy look. The woman has an intense and direct gaze,adding to the dramatic feel of the image. The backdrop is flowing silk,big silk. The overall composition blends elements of fashion and nature,creating a striking and powerful visual",
+    "A black and white portrait of a young woman with a captivating gaze. She's bundled up in a cozy black sweater,hands gently cupped near her face. The monochromatic tones highlight her delicate features and the contemplative mood of the image",
+    "Fashion photography portrait,close up portrait,(a woman of European descent is surrounded by lava rock and magma from head to neck, red magma hair, wear volcanic lava rock magma outfit coat lava rock magma fashion costume with ruffled layers"
 ]
 
 
@@ -99,7 +148,7 @@ examples = [
 
 with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
     gr.HTML("<h1><center>Flux Labs</center></h1>")
-    gr.HTML("<p><center>Model Now:  Shakker-Labs/FLUX.1-dev-LoRA-blended-realistic-illustration<br>🙇‍♂️Frequent model changes</center></p>")
+    gr.HTML("<p><center>Choose the LoRA model on the right menu</center></p>")
     with gr.Row():
         with gr.Column(scale=4):
             img = gr.Gallery(label='flux Generated Image', columns = 1, preview=True, height=600)
@@ -142,7 +191,6 @@ with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
                     maximum=MAX_SEED,
                     step=1,
                     value=-1,
-                    scale=2,
                 )
                 nums = gr.Slider(
                     label="Image Numbers",
@@ -150,25 +198,57 @@ with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
                     maximum=4,
                     step=1,
                     value=1,
-                    scale=1,
-                )  
+                )
+                lora_scale = gr.Slider(
+                    label="LoRA Scale",
+                    minimum=0.1,
+                    maximum=1.0,
+                    step=0.1,
+                    value=1.0,
+                )
+                lora_in = gr.Dropdown(
+                    ["Shakker-Labs/FLUX.1-dev-LoRA-blended-realistic-illustration", "Shakker-Labs/AWPortrait-FL", ""], 
+                    label="LoRA Model", 
+                    info="Load the LoRA",
+                    values="Shakker-Labs/FLUX.1-dev-LoRA-blended-realistic-illustration"
+                )
+                lora_add = gr.Textbox(
+                    label="Add Flux LoRA",
+                    info="Copy the HF LoRA model name here",
+                    lines=1,
+                    value="",
+                )
     gr.Examples(
         examples=examples,
-        inputs=prompt,
+        inputs=[prompt,lora_scale,lora_in],
         outputs=[img, seed],
-        fn=generate_image,
+        fn=gen,
         cache_examples="lazy",
         examples_per_page=4,
     )
+
+    gr.on(
+        triggers=[
+            prompt.submit,
+            submit.click,
+        ],
+        fn=gen,
+        inputs=[
+            prompt, 
+            width, 
+            height, 
+            scales, 
+            steps, 
+            seed, 
+            nums,
+            lora_scale,
+            lora_in,
+            lora_add
+        ],
+        outputs=[img, seed],
+        api_name="run",
+    )
     
-    sendBtn.click(fn=generate_image,
-                 inputs=[prompt, width, height, scales, steps, seed, nums],
-                 outputs=[img, seed],
-                 )
-    prompt.submit(fn=generate_image,
-                 inputs=[prompt, width, height, scales, steps, seed, nums],
-                 outputs=[img, seed],
-                 )
 
     
 demo.queue().launch()