import streamlit as st
from transformers import pipeline
from transformers import MarianMTModel, MarianTokenizer
from diffusers import StableDiffusionPipeline
from accelerate import Accelerator
import torch

# Inisialisasi accelerator
accelerator = Accelerator()

# Inisialisasi model terjemahan
model_name = "Helsinki-NLP/opus-mt-id-en"
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)

# Instalisasi Diffusion User
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)

if torch.cuda.is_available():
    # Pipe akan dipindahkan ke GPU jika tersedia
    pipe = pipe.to(accelerator.device)
    st.session_state.model_initialized = True
    st.markdown("<span style='color: green'>Model dipindahkan ke GPU.</span>", unsafe_allow_html=True)
else:
    # Memeriksa apakah menggunakan accelerate di CPU
    if accelerator.device.type == 'cpu':
        with accelerator.device:
            st.session_state.model_initialized = True
            st.markdown("<span style='color: green'>Model dijalankan di CPU dengan Accelerate.</span>", unsafe_allow_html=True)
    else:
        # Jika tidak menggunakan accelerate, tetap di CPU
        pipe = pipe.to("cpu")
        st.session_state.model_initialized = True
        st.markdown("<span style='color: red'>GPU tidak ditemukan. Model tetap di CPU.</span>", unsafe_allow_html=True)


st.title("AI Diffusion Bahasa Indonesia")
st.write("Ini adalah pembuat gambar menggunakan AI, saya menggunakan model runwayml/stable-diffusion-v1.5 dan Helsinki-NLP/opus-mt-id-e. Berkerja dengan baik, tetapi karena saya menggunakan kapasitas geratisan jadi prosesnya akan sangat lama")
text = st.text_area('Masukkan teks dalam bahasa Indonesia')

if st.button("Proses"):
    if text:
        with st.spinner("Sedang memproses..."):
            # Terjemahkan teks ke bahasa Inggris
            inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
            outputs = model.generate(**inputs)
            translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
            # Membuat Gambar
            image = pipe(translated_text).images[0]  
        
            # Tampilkan hasil
            st.image(image, caption='Hasil Diffusion User', use_column_width=True)
    else:
        st.write("Silakan masukkan teks terlebih dahulu.")