Update app.py

app.py

@@ -15,7 +15,25 @@ API_URL = "black-forest-labs/FLUX.1-schnell"
 
 mod_list = {
     "FLUX.1 Schnell": API_URL,
-    "FLUX.1 Dev": "black-forest-labs/FLUX.1-dev"
+    "FLUX.1 Schnell | Face Realism": "prithivMLmods/Canopus-LoRA-Flux-FaceRealism",
+    "FLUX.1 Schnell | Midjourney": "Jovie/Midjourney_Schnell",
+    "FLUX.1 Dev": "black-forest-labs/FLUX.1-dev",
+    "FLUX.1 Schnell | realism": "hugovntr/flux-schnell-realism",
+    "FLUX.1 Schnell | MJ v6": "strangerzonehf/Flux-Midjourney-Mix2-LoRA",
+    "FLUX.1 Dev | Flux Realism": "XLabs-AI/flux-RealismLora",
+    "FLUX.1 Redux": "black-forest-labs/FLUX.1-Redux-dev",
+    "FLUX.1 Dev | UltraRealism 2.0": "prithivMLmods/Canopus-LoRA-Flux-UltraRealism-2.0",
+    "Lumina Image 2.0": "Alpha-VLLM/Lumina-Image-2.0",
+    "FLUX.1 Dev | Vector Journey": "Shakker-Labs/FLUX.1-dev-LoRA-Vector-Journey",
+    "FLUX.1 Dev | Deep Blue": "fffiloni/deep-blue-v2",
+    "Realism Engine v1.0": "digiplay/RealismEngine_v1",
+    "Absolute Reality v1.8.1": "digiplay/AbsoluteReality_v1.8.1",
+    "FLUX.1 Dev | Midjourney Anime": "brushpenbob/flux-midjourney-anime",
+    "FLUX.1 Dev | Pencil v2": "brushpenbob/flux-pencil-v2",
+    "FLUX.1 Dev | Add Details": "Shakker-Labs/FLUX.1-dev-LoRA-add-details",
+    "FLUX.1 Dev | Pastel Anime": "Raelina/Flux-Pastel-Anime",
+    "FLUX.1 Dev | CartoonStyle": "Norod78/CartoonStyle-flux-lora",
+    "Kwai Kolors": "Kwai-Kolors/Kolors"
 }
 API_TOKEN = os.getenv("HF_READ_TOKEN")
 headers = {"Authorization": f"Bearer {API_TOKEN}"}
@@ -28,9 +46,16 @@ def add_noise(image, intensity=25):
     :param intensity: Интенсивность шума (от 0 до 255)
     :return: Изображение с шумом (PIL.Image)
     """
+    # Преобразуем изображение в массив NumPy
     img_array = np.array(image)
+    
+    # Генерируем шум
     noise = np.random.randint(-intensity, intensity, img_array.shape, dtype=np.int32)
+    
+    # Добавляем шум к изображению
     noisy_array = np.clip(img_array + noise, 0, 255).astype(np.uint8)
+    
+    # Преобразуем массив обратно в изображение
     noisy_image = Image.fromarray(noisy_array)
     return noisy_image
 
@@ -38,12 +63,19 @@ def resize_and_crop(image, target_width=768, target_height=1024):
     """
     Подгоняет изображение под размер target_width x target_height с сохранением пропорций.
     Если изображение не соответствует соотношению сторон, обрезает его по бокам.
+    :param image: Изображение (PIL.Image)
+    :param target_width: Целевая ширина
+    :param target_height: Целевая высота
+    :return: Изображение с измененным размером (PIL.Image)
     """
+    # Сохраняем исходные пропорции
     original_width, original_height = image.size
     target_ratio = target_width / target_height
     original_ratio = original_width / original_height
 
+    # Масштабируем изображение с сохранением пропорций
     if original_ratio > target_ratio:
+        # Если изображение шире, чем нужно, обрезаем по бокам
         new_height = target_height
         new_width = int(original_width * (target_height / original_height))
         resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
@@ -52,6 +84,7 @@ def resize_and_crop(image, target_width=768, target_height=1024):
         right = (new_width + target_width) / 2
         bottom = target_height
     else:
+        # Если изображение уже, чем нужно, обрезаем сверху и снизу
         new_width = target_width
         new_height = int(original_height * (target_width / original_width))
         resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
@@ -60,74 +93,82 @@ def resize_and_crop(image, target_width=768, target_height=1024):
         right = target_width
         bottom = (new_height + target_height) / 2
 
+    # Обрезаем изображение до целевого размера
     cropped_image = resized_image.crop((left, top, right, bottom))
     return cropped_image
 
 def compress_image(image, quality=50):
     """
     Сжимает изображение и понижает его качество в формате JPEG.
+    :param image: Изображение (PIL.Image)
+    :param quality: Качество JPEG (от 1 до 100)
+    :return: Сжатое изображение (PIL.Image)
     """
+    # Сохраняем изображение в буфер с пониженным качеством
     buffer = io.BytesIO()
     image.save(buffer, format="JPEG", quality=quality)
     buffer.seek(0)
+    
+    # Загружаем изображение обратно из буфера
     compressed_image = Image.open(buffer)
     return compressed_image
 
-def query(prompt, is_realistic, num_inference_steps, width, height, mod=None):
-    if not prompt:
+def query(prompt, is_realistic, num_inference_steps, width, height, mod = None):
+    if prompt == "" or prompt == None:
         return None
 
     key = random.randint(0, 999)
-    API_TOKEN = os.getenv("HF_READ_TOKEN")
+    
+    API_TOKEN = random.choice([os.getenv("HF_READ_TOKEN")])
     headers = {"Authorization": f"Bearer {API_TOKEN}"}
     
     payload = {
         "inputs": prompt,
         "seed": random.randint(1, 1000000000),
-        "parameters": {
-            "num_inference_steps": num_inference_steps,
-            "width": 1024,
-            "height": 1024
-        }
+        "parameters": {"num_inference_steps": num_inference_steps, "width": 1024, "height": 1024}
     }
 
     model = API_URL
-    if mod in mod_list:
-        model = mod_list[mod]
+    for mod_name, mod_link in mod_list.items():
+        if (mod == mod_name):
+            model = mod_link
+
+    model = "https://api-inference.huggingface.co/models/" + model;
+    
+    response = requests.post(model, 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:
-        response = requests.post(
-            f"https://api-inference.huggingface.co/models/{model}",
-            headers=headers,
-            json=payload,
-            timeout=timeout
-        )
-        response.raise_for_status()
-        
-        image = Image.open(io.BytesIO(response.content))
+        image_bytes = response.content
+        image = Image.open(io.BytesIO(image_bytes))
+
         image = resize_and_crop(image, width, height)
 
-        if is_realistic:
-            image = add_noise(image, intensity=randint(50, 100) / 10)
+        if (is_realistic == True):
+            image = add_noise(image, intensity = randint(50, 100) / 10)
             image = compress_image(image, randint(80, 90))
-            
+
             enhancer = ImageEnhance.Contrast(image)
-            image = enhancer.enhance(randint(75, 80) / 100)
+            image = enhancer.enhance(randint(75,80) / 100)
             
+            # Изменение насыщенности
             enhancer = ImageEnhance.Color(image)
-            image = enhancer.enhance(randint(80, 90) / 100)
-            
+            image = enhancer.enhance(randint(80,90) / 100)
+    
+            # Изменение экспозиции (яркости)
             enhancer = ImageEnhance.Brightness(image)
-            image = enhancer.enhance(randint(70, 100) / 100)
+            image = enhancer.enhance(randint(70,100) / 100)
 
+        
         print(f'\033[1mGeneration {key} completed!\033[0m ({prompt})')
         return image
-        
-    except requests.exceptions.RequestException as e:
-        print(f"Request error: {e}")
-        raise gr.Error(f"API Error: {str(e)}")
     except Exception as e:
-        print(f"Image processing error: {e}")
+        print(f"Error when trying to open the image: {e}")
         return None
 
 css = """
@@ -139,59 +180,27 @@ css = """
 """
 
 test_prompt = """[amateur perspective of a house],
-[on this house staying white russian female with white t-shirt, black jeans, lime sneakers and looks aside]
+[on this house staying white russian female (if you can see it) with white t-shirt (if you can see it) on her body, with black jeans (if you can see it) on her legs, with lime sneakers (if you can see it) on her feet and looks aside (if you can see it)]
 [night melancholic lighting]"""
 
 with gr.Blocks(theme='gstaff/xkcd', css=css) as app:
-    gr.Markdown("<center><h1>FLUX.1-Schnell</h1></center>")
+    gr.HTML("<center><h1>FLUX.1-Schnell</h1></center>")
     with gr.Column(elem_id="app-container"):
         with gr.Row():
             with gr.Column(elem_id="prompt-container"):
-                text_prompt = gr.Textbox(
-                    label="Prompt",
-                    placeholder="Enter a prompt here",
-                    lines=2,
-                    value=test_prompt
-                )
+                text_prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here", lines=2, elem_id="prompt-text-input", value=test_prompt)
                 is_realistic = gr.Checkbox(label="Realistic filter", value=True)
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=4
-                )
+                num_inference_steps = gr.Slider(label="Number of inference steps", minimum=1, maximum=50, step=1, value=4)
                 with gr.Row():
-                    width = gr.Slider(
-                        label="Width",
-                        minimum=128,
-                        maximum=1024,
-                        step=32,
-                        value=480
-                    )
-                    height = gr.Slider(
-                        label="Height",
-                        minimum=128,
-                        maximum=1024,
-                        step=32,
-                        value=640
-                    )
-                mod_choice = gr.Dropdown(
-                    list(mod_list.keys()),
-                    label="Model",
-                    value="FLUX.1 Schnell"
-                )
+                    width = gr.Slider(label="Width", minimum = 128, maximum = 1024, step = 32, value = 480)
+                    height = gr.Slider(label="Height", minimum = 128, maximum = 1024, step = 32, value = 640)
+                mod_choice = gr.Dropdown(list(mod_list.keys()), label="Mod")
                     
         with gr.Row():
-            text_button = gr.Button("Generate", variant='primary')
+            text_button = gr.Button("Run", variant='primary', elem_id="gen-button")
         with gr.Row():
-            image_output = gr.Image(type="pil", label="Generated Image")
-
-    text_button.click(
-        fn=query,
-        inputs=[text_prompt, is_realistic, num_inference_steps, width, height, mod_choice],
-        outputs=image_output
-    )
+            image_output = gr.Image(type="pil", label="Image Output", elem_id="gallery")
+        
+        text_button.click(query, inputs=[text_prompt, is_realistic, num_inference_steps, width, height, mod_choice], outputs=image_output)
 
-if __name__ == "__main__":
-    app.launch(show_api=True, share=False)
+app.launch(show_api=True, share=True)