Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
-import requests
-import io
+import torch
+from diffusers import DiffusionPipeline, EulerDiscreteScheduler
 import random
 import os
 from PIL import Image
@@ -8,11 +8,12 @@ from deep_translator import GoogleTranslator
 import json
 from langdetect import detect
 
-api_base = "https://api-inference.huggingface.co/models/"
-model = "cagliostrolab/animagine-xl-3.1"  # Модель Animagine XL 3.1
-timeout = 100
+model_id = "cagliostrolab/animagine-xl-3.1"
+pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="main")
+pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
+pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
-def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, gpt=False):
+def query(prompt, is_negative=False, steps=30, cfg_scale=7, seed=-1, gpt=False):
     if prompt == "" or prompt == None:
         return None
 
@@ -52,8 +53,6 @@ def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Ka
         else:
             print(f"Error: {response.status_code} - {response.text}")
 
-    API_TOKEN = random.choice([os.getenv("HF_READ_TOKEN"), os.getenv("HF_READ_TOKEN_2"), os.getenv("HF_READ_TOKEN_3"), os.getenv("HF_READ_TOKEN_4"), os.getenv("HF_READ_TOKEN_5")])
-    headers = {"Authorization": f"Bearer {API_TOKEN}"}
     language = detect(prompt)
 
     if language != 'en':
@@ -63,37 +62,13 @@ def query(prompt, is_negative=False, steps=30, cfg_scale=7, sampler="DPM++ 2M Ka
     prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
     print(f'\033[1mГенерация {key}:\033[0m {prompt}')
 
-    API_URL = model
+    if seed == -1:
+        seed = random.randint(1, 1000000000)
 
-    payload = {
-        "inputs": prompt,
-        "is_negative": is_negative,
-        "steps": steps,
-        "cfg_scale": cfg_scale,
-        "seed": seed if seed != -1 else random.randint(1, 1000000000),
-        "guidance_scale": cfg_scale,
-        "num_inference_steps": steps,
-        "negative_prompt": is_negative
-    }
-
-    response = requests.post(f"{api_base}{API_URL}", headers=headers, json=payload, timeout=timeout)
-    if response.status_code != 200:
-        print(f"Ошибка: Не удалось получить изображение. Статус ответа: {response.status_code}")
-        print(f"Содержимое ответа: {response.text}")
-        if response.status_code == 503:
-            raise gr.Error(f"{response.status_code} : The model is being loaded")
-            return None
-        raise gr.Error(f"{response.status_code}")
-        return None
-
-    try:
-        image_bytes = response.content
-        image = Image.open(io.BytesIO(image_bytes))
-        print(f'\033[1mГенерация {key} завершена!\033[0m ({prompt})')
-        return image
-    except Exception as e:
-        print(f"Ошибка при попытке открыть изображение: {e}")
-        return None
+    generator = torch.Generator("cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
+    image = pipe(prompt, negative_prompt=is_negative, guidance_scale=cfg_scale, num_inference_steps=steps, generator=generator).images[0]
+    print(f'\033[1mГенерация {key} завершена!\033[0m ({prompt})')
+    return image
 
 css = """
 * {}
@@ -113,8 +88,6 @@ with gr.Blocks(css=css) as dalle:
                     steps = gr.Slider(label="Sampling steps", value=35, minimum=1, maximum=70, step=1)
                 with gr.Row():
                     cfg = gr.Slider(label="CFG Scale", value=7, minimum=1, maximum=20, step=0.1)
-                with gr.Row():
-                    method = gr.Radio(label="Sampling method", value="Euler a", choices=["DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM"])
                 with gr.Row():
                     seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=1000000000, step=1)
                 with gr.Row():
@@ -131,6 +104,6 @@ with gr.Blocks(css=css) as dalle:
             with gr.Row():
                 image_output = gr.Image(type="pil", label="Изображение", elem_id="gallery")
 
-    text_button.click(query, inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, gpt], outputs=image_output)
+    text_button.click(query, inputs=[text_prompt, negative_prompt, steps, cfg, seed, gpt], outputs=image_output)
 
-dalle.queue(max_size=200).launch(show_api=False, share=False)
+dalle.queue(max_size=100).launch(show_api=False, share=False)