Update app.py

app.py

@@ -1,50 +1,43 @@
 import gradio as gr
-from diffusers import StableDiffusionImg2ImgPipeline
 import torch
+from diffusers import StableDiffusionImg2ImgPipeline
 from PIL import Image
 
-# مدل سبک و عمومی
-MODEL_ID = "nitrosocke/anything-mid"  # سبک و img2img مناسب
-
-# لود مدل
-pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.float16
-).to("cuda" if torch.cuda.is_available() else "cpu")  # اگر GPU باشد سریع‌تر
-
-# غیرفعال کردن فیلتر NSFW
-def dummy_checker(images, clip_input):
-    return images, [False] * len(images)
-pipe.safety_checker = dummy_checker
+# مدل سبک و سریع waifu-diffusion-lite
+model_id = "hakurei/waifu-diffusion-lite"
 
-def convert(image, prompt, strength, guidance):
-    if image is None:
-        return None
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    # تغییر اندازه عکس ورودی به 560x560
-    image = image.convert("RGB").resize((560, 560))
-
-    # اجرای مدل
-    result = pipe(
-        prompt=prompt,
-        image=image,
-        strength=strength,
-        guidance_scale=guidance,
-        num_inference_steps=15  # کم برای سرعت بالا
-    ).images[0]
-
-    return result
-
-demo = gr.Interface(
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16 if device=="cuda" else torch.float32
+).to(device)
+
+# کم کردن مراحل برای سرعت بیشتر
+pipe.scheduler.num_inference_steps = 15
+
+def convert(input_image, prompt, strength=0.7, output_resolution=560):
+    # تغییر اندازه ورودی برای سرعت
+    input_image = input_image.resize((output_resolution, output_resolution))
+    
+    # تبدیل به انیمه
+    with torch.autocast(device) if device=="cuda" else torch.no_grad():
+        result = pipe(prompt=prompt, image=input_image, strength=strength).images[0]
+    
+    # بزرگنمایی به 1024×1024 با خودکارشده (اختیاری)
+    final_image = result.resize((1024, 1024))  
+    return final_image
+
+# رابط Gradio
+gr.Interface(
     fn=convert,
     inputs=[
-        gr.Image(type="pil", label="عکس ورودی"),
-        gr.Textbox(label="پرامپت", value="anime style, colorful, high contrast"),
-        gr.Slider(0.1, 1.0, value=0.7, label="قدرت تغییر (strength)"),
-        gr.Slider(1, 15, value=7.5, label="راهنمایی مدل (guidance scale)")
+        gr.Image(type="pil"),
+        gr.Textbox(label="Prompt"),
+        gr.Slider(minimum=0.1, maximum=1.0, step=0.05, value=0.7, label="Strength"),
+        gr.Dropdown(choices=[360, 480, 560, 640, 800], value=560, label="Input Resolution")
     ],
-    outputs=gr.Image(label="خروجی"),
-    title="تبدیل عکس به انیمه سبک و سریع (560×560)"
-)
-
-demo.launch()
+    outputs=gr.Image(type="pil"),
+    title="Img2Anime Fast",
+    description="تبدیل عکس به انیمه با مدل سبک و سریع. رزولوشن ورودی قابل انتخاب است."
+).launch()