Update app.py

app.py

@@ -1,130 +1,78 @@
 from fastapi import FastAPI
-from fastapi.responses import PlainTextResponse
+from fastapi.responses import PlainTextResponse, FileResponse
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
-from deep_translator import GoogleTranslator
+from diffusers import StableDiffusionPipeline
 import torch
 import uvicorn
 import threading
 import time
 from collections import OrderedDict
+import os
 
 app = FastAPI()
 
-# 🔥 MODEL
-model_name = "Qwen/Qwen2.5-1.5B-Instruct"
+# 🔥 MODEL (small Stable Diffusion)
+model_name = "segmind/small-sd"
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(
+pipe = StableDiffusionPipeline.from_pretrained(
     model_name,
-    torch_dtype=torch.float16,
-    device_map="auto"
+    torch_dtype=torch.float16
 )
 
-MAX_HISTORY = 40
-NUM_WORKERS = 3
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = pipe.to(device)
 
-db = OrderedDict()
-queue = []
+# 🔥 оптимизации
+pipe.enable_attention_slicing()
 
-class Message(BaseModel):
-    message: str
+# если доступно
+try:
+    pipe.enable_xformers_memory_efficient_attention()
+except:
+    pass
 
 
-# 🔥 split text
-def split_text(text, max_len=100):
-    return "\n".join([text[i:i+max_len] for i in range(0, len(text), max_len)])
+MAX_HISTORY = 40
+NUM_WORKERS = 2  # меньше воркеров для GPU
 
+db = OrderedDict()
+queue = []
 
-# 🔥 очистка мусора Qwen
-def clean_output(text: str):
-    bad = [
-        "system",
-        "user",
-        "assistant",
-        "<|im_start|>",
-        "<|im_end|>",
-        "You are Qwen"
-    ]
+# папка для изображений
+IMG_DIR = "images"
+os.makedirs(IMG_DIR, exist_ok=True)
 
-    for b in bad:
-        text = text.replace(b, "")
 
-    return text.strip()
+class Message(BaseModel):
+    message: str
 
 
-# 🔥 GENERATION
+# 🔥 генерация изображения
 def generate_ai_stream(message: str):
 
-    messages = [
-        {
-            "role": "system",
-            "content": (
-                "Ты умный и точный ассистент. "
-                "Отвечай логично,кратко и понятно. "
-                "отвечай ВСЕГДА на русском."
-            )
-        },
-        {
-            "role": "user",
-            "content": message
-        }
-    ]
-
-    prompt = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
-
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
-
-    gen_kwargs = dict(
-        **inputs,
-        max_new_tokens=400,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.9,
-        streamer=streamer,
-        eos_token_id=tokenizer.eos_token_id
-    )
-
-    thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
-    thread.start()
-
-    partial = ""
-
-    # 🔥 streaming
-    for text in streamer:
-        partial += text
-
-        if message in db:
-            db[message]["reply"] = split_text(partial)
+    try:
+        image = pipe(
+            message,
+            num_inference_steps=15,
+            guidance_scale=7.5
+        ).images[0]
 
-    # 🔥 чистка
-    raw = clean_output(partial)
+        filename = f"{IMG_DIR}/img_{int(time.time()*1000)}.png"
+        image.save(filename)
 
-    # 🔥 перевод (fallback)
-    try:
-        translated = GoogleTranslator(
-            source='auto',
-            target='ru'
-        ).translate(raw)
-    except:
-        translated = raw
+        result = filename
 
-    final_text = split_text(translated) + " full generated"
+    except Exception as e:
+        result = f"error: {str(e)}"
 
     if message in db:
-        db[message]["reply"] = final_text
+        db[message]["reply"] = result
         db[message]["status"] = "done"
 
-    return final_text
+    return result
 
 
-# 🔥 worker
+# 🔥 worker очередь
 def worker():
     while True:
         if queue:
@@ -135,17 +83,17 @@ def worker():
 
             generate_ai_stream(message)
         else:
-            time.sleep(0.01)
+            time.sleep(0.05)
 
 
-# 🔥 workers
+# 🔥 запуск воркеров
 for _ in range(NUM_WORKERS):
     threading.Thread(target=worker, daemon=True).start()
 
 
 @app.get("/")
 async def root():
-    return PlainTextResponse("AI server running (Qwen2.5 1.5B Instruct)")
+    return PlainTextResponse("🎨 Image AI server running (small Stable Diffusion)")
 
 
 @app.get("/ask")
@@ -176,10 +124,19 @@ async def get(message: str):
     data = db[message]
 
     if data["status"] == "pending":
-        return PlainTextResponse(data["reply"] or "thinking...")
+        return PlainTextResponse("generating...")
 
     return PlainTextResponse(data["reply"])
 
 
+# 🔥 отдача изображения
+@app.get("/image")
+async def get_image(path: str):
+    if not os.path.exists(path):
+        return PlainTextResponse("file not found")
+
+    return FileResponse(path)
+
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)