Update app.py

app.py

@@ -1,26 +1,48 @@
 import gradio as gr
 from PIL import Image
 import torch
-from diffusers import StableDiffusionImg2ImgPipeline
+from torchvision import transforms
+from huggingface_hub import hf_hub_download
 
-# بارگذاری مدل Toonify مدرن
-pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-    "lavaman131/cartoonify",
-    torch_dtype=torch.float16
-).to("cuda" if torch.cuda.is_available() else "cpu")
+# -----------------------------
+# بارگذاری مدل CodeFormer مستقیم از Hugging Face
+# -----------------------------
+model_repo = "sczhou/CodeFormer"
+model_file = "CodeFormer.pth"  # نام مدل داخل repo
+model_path = hf_hub_download(repo_id=model_repo, filename=model_file)
 
-def toonify_image(input_image):
-    # تبدیل تصویر به سبک کارتونی
-    output = pipe(image=input_image, strength=0.8)
-    return output.images[0]
+# فرض: کلاس CodeFormer موجود باشد
+from codeformer import CodeFormer
 
-# رابط کاربری با Gradio
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = CodeFormer(pretrained=model_path, device=device)
+
+# -----------------------------
+# تابع تبدیل تصویر
+# -----------------------------
+def modern_cartoon(image):
+    image = image.convert("RGB")
+    transform = transforms.Compose([
+        transforms.Resize((512, 512)),
+        transforms.ToTensor()
+    ])
+    img_tensor = transform(image).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        output_tensor = model.enhance(img_tensor, fidelity=0.8)
+
+    output_image = transforms.ToPILImage()(output_tensor.squeeze(0).cpu())
+    return output_image
+
+# -----------------------------
+# رابط Gradio
+# -----------------------------
 iface = gr.Interface(
-    fn=toonify_image,
+    fn=modern_cartoon,
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="pil"),
-    title="Modern Toonify",
-    description="Upload your face photo → get a modern cartoon version. No prompt needed."
+    title="AnimeGAN → Modern Cartoon",
+    description="Upload your AnimeGAN output → get a modern cartoon with restored face and background."
 )
 
 iface.launch()