Update app.py

app.py

@@ -1,12 +1,13 @@
 import spaces
 import os
 import gradio as gr
-import torch
+#import torch
 import numpy as np
 import random
-from diffusers import FluxPipeline
+#from diffusers import FluxPipeline
+from huggingface_hub import AsyncInferenceClient
 from translatepy import Translator
-from huggingface_hub import hf_hub_download
+#from huggingface_hub import hf_hub_download
 import requests
 import re
 
@@ -14,7 +15,7 @@ os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 translator = Translator()
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 # Constants
-model = "black-forest-labs/FLUX.1-dev"
+basemodel = "black-forest-labs/FLUX.1-dev"
 MAX_SEED = np.iinfo(np.int32).max
 
 CSS = """
@@ -30,54 +31,27 @@ JS = """function () {
   }
 }"""
 
-if torch.cuda.is_available():
-    pipe = FluxPipeline.from_pretrained(model, torch_dtype=torch.bfloat16)
+client = AsyncInferenceClient()
 
-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):
+def enable_lora(lora_in, lora_add):
     pipe.unload_lora_weights()
     if not lora_in and not lora_add:
-        return
+        return basemodel
     else:
         if lora_add:
             lora_in = lora_add
-        url = f'https://huggingface.co/{lora_in}/tree/main'
-        lora_name = scrape_lora_link(url)
-        pipe.load_lora_weights(lora_in, weight_name=lora_name)
-        pipe.fuse_lora(lora_scale=lora_scale)   
+        return lora_in
 
-@spaces.GPU(duration=100)
+@spaces.GPU()
 def generate_image(
     prompt:str,
+    model:str,
     width:int=768,
     height:int=1024,
     scales:float=3.5,
     steps:int=24,
-    seed:int=-1,
-    nums:int=1):
+    seed:int=-1):
 
     pipe.to(device="cuda")
 
@@ -88,26 +62,29 @@ def generate_image(
     
     text = str(translator.translate(prompt, 'English'))
 
-    generator = torch.Generator().manual_seed(seed)
-    
+    #generator = torch.Generator().manual_seed(seed)
 
-    image = pipe(
+    image1 = await client.text_to_image(
         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
+        model=basemodel,
+    )
+    image2 = await client.text_to_image(
+        prompt=text,
+        height=height,
+        width=width,
+        guidance_scale=scales,
+        num_inference_steps=steps,
+        model=model,
+    )
 
-    return image, seed
+    return image1, image2, seed
 
 def gen(
     prompt:str,
-    lora_scale:float=1.0, 
     lora_in:str="", 
     lora_add:str="",
     width:int=768,
@@ -115,11 +92,10 @@ def gen(
     scales:float=3.5,
     steps:int=24,
     seed:int=-1,
-    nums:int=1,
     progress=gr.Progress(track_tqdm=True)
 ):
-    enable_lora(lora_scale, lora_in, lora_add)
-    return generate_image(prompt,width,height,scales,steps,seed,nums)
+    model = enable_lora(lora_in, lora_add)
+    return generate_image(prompt,model,width,height,scales,steps,seed)
      
 
 
@@ -151,11 +127,13 @@ examples = [
 # Gradio Interface
 
 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>")
+    gr.HTML("<h1><center>Flux Labs(vs LoRA)</center></h1>")
+    gr.HTML("<p><center>Choose the LoRA model on the menu</center></p>")
     with gr.Row():
         with gr.Column(scale=4):
-            img = gr.Gallery(label='flux Generated Image', columns = 1, preview=True, height=600)
+            with gr.Row():
+                img1 = gr.Gallery(label='flux Generated Image', columns = 1, preview=True, height=600)
+                img2 = gr.Gallery(label='LoRA 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')
@@ -196,20 +174,6 @@ with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
                     step=1,
                     value=-1,
                 )
-                nums = gr.Slider(
-                    label="Image Numbers",
-                    minimum=1,
-                    maximum=4,
-                    step=1,
-                    value=1,
-                )
-                lora_scale = gr.Slider(
-                    label="LoRA Scale",
-                    minimum=0.1,
-                    maximum=1.0,
-                    step=0.1,
-                    value=1.0,
-                )
                 lora_in = gr.Dropdown(
                     choices=["Shakker-Labs/FLUX.1-dev-LoRA-blended-realistic-illustration", "Shakker-Labs/AWPortrait-FL",""], 
                     label="LoRA Model", 
@@ -223,8 +187,8 @@ with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
                 )
     gr.Examples(
         examples=examples,
-        inputs=[prompt,lora_scale,lora_in],
-        outputs=[img, seed],
+        inputs=[prompt,lora_in],
+        outputs=[img1, img2, seed],
         fn=gen,
         cache_examples="lazy",
         examples_per_page=4,
@@ -238,17 +202,15 @@ with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
         fn=gen,
         inputs=[
             prompt,
-            lora_scale,
             lora_in,
             lora_add,
             width, 
             height, 
             scales, 
             steps, 
-            seed, 
-            nums
+            seed
         ],
-        outputs=[img, seed],
+        outputs=[img1, img2, seed],
         api_name="run",
     )