Update app.py

app.py

@@ -1,5 +1,5 @@
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
 import torch
 import gradio as gr
 
@@ -9,50 +9,18 @@ MODEL_ID = "microsoft/bitnet-b1.58-2B-4T-gguf"  # Инструктивная в
 # --- Ленивая загрузка модели ---
 model_loaded = False
 
-@spaces.GPU(duration=120)
-def load_model():
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        device_map="auto",
-        torch_dtype=torch.bfloat16,
-        trust_remote_code=True
-    )
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-    return model, tokenizer
+import spaces
+from diffusers import DiffusionPipeline
 
-# --- Генерация ---
-@spaces.GPU
-def generate_response(message: str, history=None):
-    try:
-        model, tokenizer = load_model()
-        
-        # Форматируем сообщение с историей (если нужно)
-        chat_history = history if history else []
-        prompt = "\n".join([f"User: {h[0]}\nAssistant: {h[1]}" for h in chat_history] + [f"User: {message}\nAssistant:"])
-        
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-        
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=128,  # Сильно уменьшаем для ZeroGPU
-            temperature=0.7,
-            do_sample=True
-        )
-        
-        return tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
-    
-    except Exception as e:
-        return f"Ошибка: {str(e)}"
+pipe = DiffusionPipeline.from_pretrained(MODEL_ID)
+pipe.to('cuda')
 
-# --- Интерфейс ---
-with gr.Blocks() as demo:
-    gr.Markdown("## 🤖 Mistral-7B Chat (ZeroGPU)")
-    
-    chat = gr.ChatInterface(
-        generate_response,
-        examples=["Привет!", "Объясни квантовую физику просто"],
-        title="Mistral-7B"
-    )
+@spaces.GPU
+def generate(prompt):
+    return pipe(prompt).images
 
-if __name__ == "__main__":
-    demo.launch()
+gr.Interface(
+    fn=generate,
+    inputs=gr.Text(),
+    outputs=gr.Gallery(),
+).launch()