Update app.py

app.py

@@ -1,28 +1,20 @@
-import subprocess
-command = 'pip install git+https://github.com/snekkenull/diffusers.git'
-subprocess.run(command, shell=True)
-
+import spaces
 import os
 import gradio as gr
 import torch
 import numpy as np
 import random
-from diffusers import StableDiffusion3Pipeline, AutoencoderKL, SD3Transformer2DModel, StableDiffusion3Img2ImgPipeline, FlowMatchEulerDiscreteScheduler
-import spaces
-from diffusers.utils import load_image
-from PIL import Image
-import requests
-import transformers
-from transformers import AutoTokenizer, T5EncoderModel
+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()
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 # Constants
-model = "stabilityai/stable-diffusion-3-medium"
-repo= "stabilityai/stable-diffusion-3-medium-diffusers"
+model = "black-forest-labs/FLUX.1-dev"
 MAX_SEED = np.iinfo(np.int32).max
 
 CSS = """
@@ -38,143 +30,108 @@ JS = """function () {
   }
 }"""
 
+if torch.cuda.is_available():
+    pipe = FluxPipeline.from_pretrained(model, torch_dtype=torch.bfloat16)
 
-vae = AutoencoderKL.from_pretrained(
-    repo,
-    subfolder="vae",
-    torch_dtype=torch.float16,
- )
-
-transformer = SD3Transformer2DModel.from_pretrained(
-    repo, 
-    subfolder="transformer",
-    torch_dtype=torch.float16,
-)
-
-
-# text_encoder_3 = T5EncoderModel.from_pretrained(
-#     repo, 
-#     subfolder="text_encoder_3", 
-# )
-
-# tokenizer_3 = AutoTokenizer.from_pretrained(
-#     repo, 
-#     subfolder="tokenizer_3",
-#     torch_dtype=torch.float16,
-# )
 
+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
 
-# Ensure model and scheduler are initialized in GPU-enabled function
-if torch.cuda.is_available():
-    pipe = StableDiffusion3Pipeline.from_pretrained(
-        repo,
-        vae=vae,
-        transformer=transformer,
-        torch_dtype=torch.float16).to("cuda")
-    pipe2 = StableDiffusion3Img2ImgPipeline.from_pretrained(
-        repo,
-        vae=vae,
-        transformer=transformer,
-        torch_dtype=torch.float16).to("cuda")
+        # 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
 
-pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config)
-pipe2.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe2.scheduler.config)
+    except requests.RequestException as e:
+        print(f"An error occurred while fetching the URL: {e}")
+        return None
 
-print(pipe.tokenizer_max_length)
+def enable_lora(lora_scale,lora_add):
+    if not lora_add:
+        gradio.Info("No Lora Loaded, Use basemodel", duration=5)
+    else:
+        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)
+        gradio.Info(f"{lora_add} Loaded", duration=5)
 
-# Function 
 @spaces.GPU()
 def generate_image(
-    prompt,
-    negative="low quality",
-    width=1024,
-    height=1024,
-    scales=5,
-    steps=30,
-    strength=0.7,
-    seed: int =-1,
-    nums=1,
-    progress=gr.Progress(track_tqdm=True)):
+    prompt:str,
+    lora_word:str,
+    width:int=768,
+    height:int=1024,
+    scales:float=3.5,
+    steps:int=24,
+    seed:int=-1,
+    nums:int=1):
+
+    pipe.to(device="cuda")
 
     if seed == -1:
         seed = random.randint(0, MAX_SEED)
     seed = int(seed)
     print(f'prompt:{prompt}')
     
-    text = str(translator.translate(prompt['text'], 'English'))
+    text = str(translator.translate(prompt, 'English')) + "," + lora_word
 
-    
-    if prompt['files']:
-        #images = Image.open(prompt['files'][-1]).convert('RGB')
-        init_image = load_image(prompt['files'][-1]).resize((height, width))
-    else:
-        init_image = None
     generator = torch.Generator().manual_seed(seed)
     
 
-    if init_image:
-        image = pipe2(
-            prompt=text,
-            image=init_image,
-            negative_prompt=negative, 
-            guidance_scale=scales,
-            num_inference_steps=steps,
-            strength=strength,
-            generator = generator,
-            num_images_per_prompt = nums,
-        ).images
-    else:
-        image = pipe(
-            prompt=text,
-            negative_prompt=negative, 
-            width=width,
-            height=height,
-            guidance_scale=scales,
-            num_inference_steps=steps,
-            generator = generator,
-            num_images_per_prompt = nums,
-        ).images
-
-    print(image)
-    print(seed)
-    return image, seed
-
+    image = pipe(
+        prompt=text,
+        height=height,
+        width=width,
+        guidance_scale=scales,
+        output_type="pil",
+        num_inference_steps=steps,
+        max_sequence_length=512,
+        num_images_per_prompt=nums,
+        generator=generator,
+    ).images
 
+    return image, seed
+ 
 examples = [
-        [{"text": "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", "files": []}],
-        [{"text": "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\".", "files": []}],
-        [{"text": "photo of a dog and a cat both standing on a red box, with a blue ball in the middle with a parrot standing on top of the ball. The box has the text \"SD3\"", "files": []}],
-		[{"text": "selfie photo of a wizard with long beard and purple robes, he is apparently in the middle of Tokyo. Probably taken from a phone.", "files": []}],
-		[{"text": "A vibrant street wall covered in colorful graffiti, the centerpiece spells \"SD3 MEDIUM\", in a storm of colors", "files": []}],
-		[{"text":  "photo of a young woman with long, wavy brown hair tied in a bun and glasses. She has a fair complexion and is wearing subtle makeup, emphasizing her eyes and lips. She is dressed in a black top. The background appears to be an urban setting with a building facade, and the sunlight casts a warm glow on her face.", "files": []}],
-		[{"text": "anime art of a steampunk inventor in their workshop, surrounded by gears, gadgets, and steam. He is holding a blue potion and a red potion, one in each hand", "files": []}],
-		[{"text": "photo of picturesque scene of a road surrounded by lush green trees and shrubs. The road is wide and smooth, leading into the distance. On the right side of the road, there's a blue sports car parked with the license plate spelling \"SD32B\". The sky above is partly cloudy, suggesting a pleasant day. The trees have a mix of green and brown foliage. There are no people visible in the image. The overall composition is balanced, with the car serving as a focal point.", "files": []}],
-		[{"text": "photo of young man in a black suit, white shirt, and black tie. He has a neatly styled haircut and is looking directly at the camera with a neutral expression. The background consists of a textured wall with horizontal lines. The photograph is in black and white, emphasizing contrasts and shadows. The man appears to be in his late twenties or early thirties, with fair skin and short, dark hair.", "files": []}],
-		[{"text": "photo of a woman on the beach, shot from above. She is facing the sea, while wearing a white dress. She has long blonde hair", "files": []}],
+    ["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.",0.9,"Shakker-Labs/AWPortrait-FL",""],
+    ["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",0.9,"Shakker-Labs/AWPortrait-FL".""],
+    ["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",0.9,"Shakker-Labs/AWPortrait-FL",""],
+    ["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",0.9,"Shakker-Labs/AWPortrait-FL",""]
 ]
 
 
-
 # Gradio Interface
 
-with gr.Blocks(css=CSS, js=JS, theme="soft") as demo:
-    gr.HTML("<h1><center>SD3M🐉T5</center></h1>")
-    gr.HTML("<p><center><a href='https://huggingface.co/stabilityai/stable-diffusion-3-medium'>sd3m</a> text/image-to-image generation<br><b>Update</b>: fix diffuser to support 512 token</center></p>")
+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>Choose the LoRA model on the right menu</center></p>")
     with gr.Row():
         with gr.Column(scale=4):
-            img = gr.Gallery(label='SD3M Generated Image', columns = 1, preview=True, height=600)
-            prompt = gr.MultimodalTextbox(label='Enter Your Prompt (Multi-Languages)', interactive=True, placeholder="Enter prompt, add one image.", file_types=['image'])
+            img = gr.Gallery(label='flux Generated Image', columns = 1, preview=True, height=600)
+            with gr.Row():
+                prompt = gr.Textbox(label='Enter Your Prompt (Multi-Languages)', placeholder="Enter prompt...", scale=6)
+                sendBtn = gr.Button(scale=1, variant='primary')
         with gr.Accordion("Advanced Options", open=True):
             with gr.Column(scale=1):
-                negative = gr.Textbox(label="Negative prompt", value="low quality, ugly, blurry, poor face, bad anatomy")
                 width = gr.Slider(
                     label="Width",
                     minimum=512,
                     maximum=1280,
                     step=8,
-                    value=1024,
+                    value=768,
                 )
                 height = gr.Slider(
                     label="Height",
@@ -188,29 +145,21 @@ with gr.Blocks(css=CSS, js=JS, theme="soft") as demo:
                     minimum=3.5,
                     maximum=7,
                     step=0.1,
-                    value=5,
+                    value=3.5,
                 )
                 steps = gr.Slider(
                     label="Steps",
                     minimum=1,
-                    maximum=50,
+                    maximum=100,
                     step=1,
-                    value=30,
-                )
-                strength = gr.Slider(
-                    label="Strength",
-                    minimum=0.0,
-                    maximum=1.0,
-                    step=0.1,
-                    value=0.7,
+                    value=24,
                 )
                 seed = gr.Slider(
-                    label="Seed (-1 Random)",
+                    label="Seeds",
                     minimum=-1,
                     maximum=MAX_SEED,
                     step=1,
                     value=-1,
-                    scale=2,
                 )
                 nums = gr.Slider(
                     label="Image Numbers",
@@ -218,21 +167,57 @@ with gr.Blocks(css=CSS, js=JS, theme="soft") 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_add = gr.Textbox(
+                    label="Add Flux LoRA",
+                    info="Copy the HF LoRA model name here",
+                    lines=1,
+                    value="Shakker-Labs/AWPortrait-FL",
+                )
+                lora_word = gr.Textbox(
+                    label="Add Flux LoRA Trigger Word",
+                    info="Add the Trigger Word",
+                    lines=1,
+                    value="",
+                )
+                load_lora = gr.Button(value="Load LoRA", variant='secondary')
+                
     gr.Examples(
         examples=examples,
-        inputs=prompt,
-        outputs=[img, seed],
-        fn=generate_image,
-        cache_examples="lazy",
+        inputs=[prompt,lora_scale,lora_add,lora_word],
+        cache_examples=False,
         examples_per_page=4,
     )
 
-    prompt.submit(fn=generate_image,
-                 inputs=[prompt, negative, width, height, scales, steps, strength, seed, nums],
-                 outputs=[img, seed],
-                 )
+    load_lora.click(fn=enable_lora, inputs=[lora_scale,lora_add])
+    
+    gr.on(
+        triggers=[
+            prompt.submit,
+            sendBtn.click,
+        ],
+        fn=generate_image,
+        inputs=[
+            prompt,
+            lora_word,
+            width, 
+            height, 
+            scales, 
+            steps, 
+            seed, 
+            nums
+        ],
+        outputs=[img, seed],
+        api_name="run",
+    )
+    
 
     
 demo.queue().launch()