Upload app.py

app.py

@@ -0,0 +1,65 @@
+import streamlit as st
+import torch
+from torchvision import transforms
+from PIL import Image
+from model import CycleGAN # Подтягиваем архитектуру из вашего model.py
+
+st.set_page_config(page_title="CycleGAN Demo", layout="wide")
+st.title("🔄 CycleGAN: Трансфер стилей")
+
+# Кешируем загрузку, чтобы не грузить 200 МБ при каждом нажатии
+@st.cache_resource
+def load_model():
+    # На бесплатном сервере HF нет видеокарт, заставляем работать на CPU
+    device = torch.device("cpu") 
+    model = CycleGAN()
+    # Загружаем ваш файл (впишите точное имя файла, который вы загрузили!)
+    checkpoint = torch.load("cycle_gan_fixed.pt", map_location=device, weights_only=False)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval()
+    return model, device
+
+model, device = load_model()
+
+# Трансформации (размер должен быть таким же, как при обучении, например 128)
+IMG_SIZE = 128 
+transform = transforms.Compose([
+    transforms.Resize((IMG_SIZE, IMG_SIZE), Image.BICUBIC),
+    transforms.ToTensor(),
+    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+])
+
+def de_normalize(tensor):
+    tensor = tensor.cpu().squeeze(0)
+    tensor = tensor * 0.5 + 0.5
+    tensor = torch.clamp(tensor, 0, 1)
+    return tensor.permute(1, 2, 0).numpy()
+
+# Интерфейс
+col1, col2 = st.columns(2)
+
+with col1:
+    st.header("Домен A ➡️ Домен B")
+    file_a = st.file_uploader("Загрузить фото A", type=["jpg", "png", "jpeg"], key="a")
+    if file_a:
+        img_a = Image.open(file_a).convert("RGB")
+        st.image(img_a, caption="Оригинал")
+        if st.button("Преобразовать", key="btn_a"):
+            with st.spinner("Генерация..."):
+                tensor = transform(img_a).unsqueeze(0)
+                with torch.no_grad():
+                    res = model.G_A2B(tensor) # Перевод из A в B
+                st.image(de_normalize(res), caption="Результат")
+
+with col2:
+    st.header("Домен B ➡️ Домен A")
+    file_b = st.file_uploader("Загрузить фото B", type=["jpg", "png", "jpeg"], key="b")
+    if file_b:
+        img_b = Image.open(file_b).convert("RGB")
+        st.image(img_b, caption="Оригинал")
+        if st.button("Преобразовать", key="btn_b"):
+            with st.spinner("Генерация..."):
+                tensor = transform(img_b).unsqueeze(0)
+                with torch.no_grad():
+                    res = model.G_B2A(tensor) # Перевод из B в A
+                st.image(de_normalize(res), caption="Результат")