Update app.py

app.py

@@ -1,75 +1,41 @@
-import gradio as gr
-from PIL import Image
 import torch
 from torchvision import transforms
-import torch.nn as nn
-import torch.nn.functional as F
-
-# ---------------------
-# Toonify Lite / CartoonGAN سبک برای CPU
-# ---------------------
-class ToonifyLite(nn.Module):
-    def __init__(self):
-        super().__init__()
-        # Encoder ساده
-        self.encoder = nn.Sequential(
-            nn.Conv2d(3, 32, 3, 1, 1), nn.ReLU(),
-            nn.Conv2d(32, 64, 3, 2, 1), nn.ReLU(),
-            nn.Conv2d(64, 128, 3, 2, 1), nn.ReLU(),
-            nn.Conv2d(128, 256, 3, 2, 1), nn.ReLU()
-        )
-        # Decoder سبک
-        self.decoder = nn.Sequential(
-            nn.Conv2d(256, 128, 3, 1, 1), nn.ReLU(),
-            nn.Upsample(scale_factor=2, mode="nearest"),
-            nn.Conv2d(128, 64, 3, 1, 1), nn.ReLU(),
-            nn.Upsample(scale_factor=2, mode="nearest"),
-            nn.Conv2d(64, 32, 3, 1, 1), nn.ReLU(),
-            nn.Upsample(scale_factor=2, mode="nearest"),
-            nn.Conv2d(32, 3, 3, 1, 1)
-        )
-
-    def forward(self, x):
-        feat = self.encoder(x)
-        out = self.decoder(feat)
-        return out
+from PIL import Image
+import requests
+from io import BytesIO
+import gradio as gr
+from config import MODEL_URL, DEVICE, IMAGE_SIZE
 
-# ---------------------
-# مدل آماده و روی CPU
-# ---------------------
-device = torch.device("cpu")
-model = ToonifyLite().to(device)
+# بارگذاری مدل مستقیم از HF
+model_path = MODEL_URL
+model = torch.hub.load('justinpinkney/toonify', 'toonify')
 model.eval()
-
-# ---------------------
-# تابع تبدیل تصویر
-# ---------------------
-def toonify_image(input_image):
-    transform = transforms.Compose([
-        transforms.Resize((512, 512)),
-        transforms.ToTensor()
+model.to(DEVICE)
+
+# تابع تبدیل عکس به Toonify
+def convert_to_toon(image: Image.Image) -> Image.Image:
+    # تغییر اندازه و تبدیل به tensor
+    preprocess = transforms.Compose([
+        transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
+        transforms.ToTensor(),
     ])
-    img = transform(input_image).unsqueeze(0).to(device)
+    img_tensor = preprocess(image).unsqueeze(0).to(DEVICE)
 
     with torch.no_grad():
-        output = model(img)
-
-    # نرمال سازی خروجی
-    output = output.squeeze(0).cpu()
-    output = (output - output.min()) / (output.max() - output.min() + 1e-5)
-    output = output.clamp(0, 1)
+        output_tensor = model(img_tensor)
 
-    return transforms.ToPILImage()(output)
+    # تبدیل خروجی به تصویر PIL
+    output_image = output_tensor.squeeze().cpu()
+    output_image = transforms.ToPILImage()(output_image)
+    return output_image
 
-# ---------------------
-# رابط کاربری Gradio
-# ---------------------
+# رابط Gradio
 iface = gr.Interface(
-    fn=toonify_image,
-    inputs=gr.Image(type="pil", label="Input Image"),
-    outputs=gr.Image(type="pil", label="Toonified Output"),
-    title="Toonify Lite - CPU",
-    description="تصویر خود را آپلود کنید و خروجی کارتونی / انیمه دریافت کنید. جزئیات چهره حفظ می‌شود."
+    fn=convert_to_toon,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="Toonify Face Converter",
+    description="تبدیل عکس شما به کارتون مدرن با حفظ جزئیات چهره، بدون نیاز به پرامپت!"
 )
 
-iface.launch(server_name="0.0.0.0", server_port=7860)
+iface.launch()