Update app.py

app.py

@@ -1,17 +1,25 @@
 import gradio as gr
-import requests
-import io
+import numpy as np
 import random
-import os
-import time
-from PIL import Image
-from deep_translator import GoogleTranslator
-import json
+import spaces
+import torch
+from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
+from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
+from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
 
 
-API_TOKEN = os.getenv("HF_READ_TOKEN")
-headers = {"Authorization": f"Bearer {API_TOKEN}"}
-timeout = 100
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
+good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=taef1).to(device)
+torch.cuda.empty_cache()
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 2048
+
+pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
 article_text = """
 <div style="text-align: center;">
@@ -27,57 +35,104 @@ article_text = """
 </div>
 """
 
-def query(lora_id, prompt, steps=28, cfg_scale=3.5, randomize_seed=True, seed=-1, width=1024, height=1024):
-    if prompt == "" or prompt == None:
-        return None
-
-    if lora_id.strip() == "" or lora_id == None:
-        lora_id = "black-forest-labs/FLUX.1-dev" 
-
-    key = random.randint(0, 999)
+@spaces.GPU()
+def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, lora_id=None, lora_scale=0.95, progress=gr.Progress(track_tqdm=True)):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator().manual_seed(seed)
 
-    API_URL = "https://api-inference.huggingface.co/models/"+ lora_id.strip()
     
-    API_TOKEN = random.choice([os.getenv("HF_READ_TOKEN")])
-    headers = {"Authorization": f"Bearer {API_TOKEN}"}
+    # for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+    #         prompt=prompt,
+    #         guidance_scale=guidance_scale,
+    #         num_inference_steps=num_inference_steps,
+    #         width=width,
+    #         height=height,
+    #         generator=generator,
+    #         output_type="pil",
+    #         good_vae=good_vae,
+    #     ):
+    #         yield img, seed
+
+    # Handle LoRA loading
+    # Load LoRA weights and prepare joint_attention_kwargs
+    if lora_id:
+        pipe.unload_lora_weights()
+        pipe.load_lora_weights(lora_id)
+        joint_attention_kwargs = {"scale": lora_scale}
+    else:
+        joint_attention_kwargs = None
     
-    # prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
-    # print(f'\033[1mGeneration {key} translation:\033[0m {prompt}')
+    try:
+        # Call the custom pipeline function with the correct keyword argument
+        for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator,
+            output_type="pil",
+            good_vae=good_vae,  # Assuming good_vae is defined elsewhere
+            joint_attention_kwargs=joint_attention_kwargs,  # Fixed parameter name
+        ):
+            yield img, seed
+    finally:
+        # Unload LoRA weights if they were loaded
+        if lora_id:
+            pipe.unload_lora_weights()
+
+# def query(lora_id, prompt, steps=28, cfg_scale=3.5, randomize_seed=True, seed=-1, width=1024, height=1024):
+#     if prompt == "" or prompt == None:
+#         return None
+
+#     if lora_id.strip() == "" or lora_id == None:
+#         lora_id = "black-forest-labs/FLUX.1-dev" 
+
+#     key = random.randint(0, 999)
+
+#     API_URL = "https://api-inference.huggingface.co/models/"+ lora_id.strip()
+    
+#     API_TOKEN = random.choice([os.getenv("HF_READ_TOKEN")])
+#     headers = {"Authorization": f"Bearer {API_TOKEN}"}
+    
+#     # prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
+#     # print(f'\033[1mGeneration {key} translation:\033[0m {prompt}')
 
-    prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
-    # print(f'\033[1mGeneration {key}:\033[0m {prompt}')
+#     prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
+#     # print(f'\033[1mGeneration {key}:\033[0m {prompt}')
 
-    # If seed is -1, generate a random seed and use it
-    if randomize_seed:
-        seed = random.randint(1, 4294967296)
+#     # If seed is -1, generate a random seed and use it
+#     if randomize_seed:
+#         seed = random.randint(1, 4294967296)
     
-    payload = {
-        "inputs": prompt,
-        "steps": steps,
-        "cfg_scale": cfg_scale,
-        "seed": seed,
-        "parameters": {
-            "width": width,  # Pass the width to the API
-            "height": height  # Pass the height to the API
-        }
-    }
-
-    response = requests.post(API_URL, headers=headers, json=payload, timeout=timeout)
-    if response.status_code != 200:
-        print(f"Error: Failed to get image. Response status: {response.status_code}")
-        print(f"Response content: {response.text}")
-        if response.status_code == 503:
-            raise gr.Error(f"{response.status_code} : The model is being loaded")
-        raise gr.Error(f"{response.status_code}")
+#     payload = {
+#         "inputs": prompt,
+#         "steps": steps,
+#         "cfg_scale": cfg_scale,
+#         "seed": seed,
+#         "parameters": {
+#             "width": width,  # Pass the width to the API
+#             "height": height  # Pass the height to the API
+#         }
+#     }
+
+#     response = requests.post(API_URL, headers=headers, json=payload, timeout=timeout)
+#     if response.status_code != 200:
+#         print(f"Error: Failed to get image. Response status: {response.status_code}")
+#         print(f"Response content: {response.text}")
+#         if response.status_code == 503:
+#             raise gr.Error(f"{response.status_code} : The model is being loaded")
+#         raise gr.Error(f"{response.status_code}")
     
-    try:
-        image_bytes = response.content
-        image = Image.open(io.BytesIO(image_bytes))
-        print(f'\033[1mGeneration {key} completed!\033[0m ({prompt})')
-        return image, seed, seed
-    except Exception as e:
-        print(f"Error when trying to open the image: {e}")
-        return None
+#     try:
+#         image_bytes = response.content
+#         image = Image.open(io.BytesIO(image_bytes))
+#         print(f'\033[1mGeneration {key} completed!\033[0m ({prompt})')
+#         return image, seed, seed
+#     except Exception as e:
+#         print(f"Error when trying to open the image: {e}")
+#         return None
 
 
 examples = [
@@ -114,13 +169,20 @@ with gr.Blocks(css=css) as app:
                 with gr.Row():
                     with gr.Accordion("Advanced Settings", open=False):
                         with gr.Row():
+                            lora_scale = gr.Slider(
+                                label="LoRA Scale",
+                                minimum=0,
+                                maximum=2,
+                                step=0.01,
+                                value=0.95,
+                            )
                             width = gr.Slider(label="Width", value=1024, minimum=64, maximum=1216, step=8)
                             height = gr.Slider(label="Height", value=1024, minimum=64, maximum=1216, step=8)
                         seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=4294967296, step=1)
                         randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                         with gr.Row():
-                            steps = gr.Slider(label="Sampling steps", value=28, minimum=1, maximum=100, step=1)
-                            cfg = gr.Slider(label="CFG Scale", value=3.5, minimum=1, maximum=20, step=0.5)
+                            steps = gr.Slider(label="Inference steps steps", value=28, minimum=1, maximum=100, step=1)
+                            cfg = gr.Slider(label="Guidance Scale", value=3.5, minimum=1, maximum=20, step=0.5)
                         # method = gr.Radio(label="Sampling method", value="DPM++ 2M Karras", choices=["DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM"])
 
                 with gr.Row():
@@ -140,6 +202,7 @@ with gr.Blocks(css=css) as app:
             )
 
         
-        text_button.click(query, inputs=[custom_lora, text_prompt, steps, cfg, randomize_seed, seed, width, height], outputs=[image_output,seed_output, seed])
+        # text_button.click(query, inputs=[custom_lora, text_prompt, steps, cfg, randomize_seed, seed, width, height], outputs=[image_output,seed_output, seed])
+        text_button.click(infer, inputs=[text_prompt, seed, randomize_seed, width, height, cfg, steps, custom_lora, lora_scale], outputs=[image_output,seed_output, seed])
 
-app.launch(show_api=False, share=False)
+app.launch()