modify

app.py

@@ -8,22 +8,28 @@ import os
 import torch
 from transformers import utils
 
-# 清理缓存
-utils.move_cache()
-
 # 从环境变量中获取访问令牌
 hf_token = os.getenv("hf_token")
 
 # 使用令牌进行登录（如果不用Hugging Face的模型，可以跳过这一步）
 if hf_token:
     login(token=hf_token)
-
-# 加载预训练的图像到图像的 Stable Diffusion 模型
-pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
-)
-
-pipe.to("cuda")  # 使用 CPU 来处理，如果有 GPU，可修改为 "cuda"
+else:
+    print("Hugging Face token not found. Please set the 'hf_token' environment variable.")
+
+# 检查是否有GPU可用，并选择相应的设备和数据类型
+if torch.cuda.is_available():
+    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
+    )
+    pipe.to("cuda")
+    print("Using CUDA")
+else:
+    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32
+    )
+    pipe.to("cpu")
+    print("CUDA not available, using CPU")
 
 # 定义 Gradio 接口中调用的函数
 def image_to_image(input_image, prompt):
@@ -34,12 +40,16 @@ def image_to_image(input_image, prompt):
     # 将用户上传的输入图像转化为 RGB 格式
     input_image = input_image.convert("RGB")
     
-    # 使用预训练模型生成新图像
-    generated_image = pipe(
-        prompt=prompt, image=input_image, strength=0.75, guidance_scale=7.5
-    ).images[0]
+    try:
+        # 使用预训练模型生成新图像
+        generated_image = pipe(
+            prompt=prompt, image=input_image, strength=0.75, guidance_scale=7.5
+        ).images[0]
     
-    return generated_image  # 返回生成的图像
+        return generated_image  # 返回生成的图像
+    except Exception as e:
+        print(f"Error during image generation: {e}")
+        return None
 
 # 创建 Gradio 接口
 demo = gr.Interface(
@@ -49,4 +59,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch(share=True)  # 如果需要公开链接，可以设置 share=True
+    demo.launch(share=False)  # 如果需要公开链接，可以设置 share=True