Update app.py

app.py

@@ -12,10 +12,8 @@ from transformers import CLIPTextModel, CLIPTokenizer
 from briarmbg import BriaRMBG
 from torch.hub import download_url_to_file
 
-# تنظیم متغیرهای محیطی
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-os.environ['CUDA_HOME'] = '/usr/local/cuda'
-os.environ['PATH'] = os.environ['PATH'] + ':/usr/local/cuda/bin'
+# نصب کتابخانه‌های مورد نیاز در صورت اجرا روی Hugging Face Spaces
+os.system("pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121")
 
 # بررسی دسترسی به GPU
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -56,7 +54,8 @@ scheduler = DPMSolverMultistepScheduler(
 )
 
 # ساخت لوله‌های پردازش تصویر
-t2i_pipe = StableDiffusionPipeline(
+t2i_pipe = StableDiffusionPipeline.from_pretrained(
+    sd15_name,
     vae=vae,
     text_encoder=text_encoder,
     tokenizer=tokenizer,
@@ -65,11 +64,11 @@ t2i_pipe = StableDiffusionPipeline(
     safety_checker=None,
     requires_safety_checker=False,
     feature_extractor=None,
-    image_encoder=None,
-    torch_dtype=dtype
+    image_encoder=None
 ).to(device)
 
-i2i_pipe = StableDiffusionImg2ImgPipeline(
+i2i_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+    sd15_name,
     vae=vae,
     text_encoder=text_encoder,
     tokenizer=tokenizer,
@@ -78,8 +77,7 @@ i2i_pipe = StableDiffusionImg2ImgPipeline(
     safety_checker=None,
     requires_safety_checker=False,
     feature_extractor=None,
-    image_encoder=None,
-    torch_dtype=dtype
+    image_encoder=None
 ).to(device)
 
 # پاک کردن کش حافظه GPU برای جلوگیری از خطای Out Of Memory
@@ -116,4 +114,4 @@ with block:
     
     submit.click(fn=process_image, inputs=[prompt, width, height, steps, guidance], outputs=result_image)
 
-block.launch(server_name='0.0.0.0')
+block.launch(server_name='0.0.0.0')