import gradio as gr
from transformers import pipeline
from diffusers import StableDiffusionPipeline
from PIL import Image
import torch

# Загрузка модели Qwen/QVQ-72B-preview (пример для текстовой модели)
def load_text_model():
    text_pipeline = pipeline("text-generation", model="Qwen/QVQ-72B-preview")
    return text_pipeline

# Загрузка модели fffiloni/AniDoc (пример для модели анализа изображений)
def load_image_model():
    image_pipeline = StableDiffusionPipeline.from_pretrained("fffiloni/AniDoc")
    return image_pipeline

# Функция для генерации текста
def generate_text(prompt, text_pipeline):
    result = text_pipeline(prompt, max_length=50)
    return result[0]['generated_text']

# Функция для анализа изображения
def analyze_image(image, image_pipeline):
    # Пример: генерация изображения на основе входного изображения
    prompt = "Anime style"
    output_image = image_pipeline(prompt, image=image).images[0]
    return output_image

# Загрузка моделей
text_pipeline = load_text_model()
image_pipeline = load_image_model()

# Создание Gradio интерфейса
with gr.Blocks() as demo:
    with gr.Tab("Text Generation"):
        text_input = gr.Textbox(label="Input Prompt")
        text_output = gr.Textbox(label="Generated Text")
        text_button = gr.Button("Generate Text")
        text_button.click(fn=generate_text, inputs=[text_input, gr.State(text_pipeline)], outputs=text_output)

    with gr.Tab("Image Analysis"):
        image_input = gr.Image(label="Input Image", type="pil")
        image_output = gr.Image(label="Output Image")
        image_button = gr.Button("Analyze Image")
        image_button.click(fn=analyze_image, inputs=[image_input, gr.State(image_pipeline)], outputs=image_output)

# Запуск интерфейса
demo.launch()