lkspace

app.py

@@ -12,11 +12,12 @@ login(token=os.getenv("HF_TOKEN"))
 
 # 4. Load private model using your token
 model_id = "Lookingsoft-team/Text_to_Image_Diffusion"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 pipe = StableDiffusionPipeline.from_pretrained(
     model_id,
     torch_dtype=torch.float16  # smaller memory
-).to("cuda")
+).to(device)
 
 def generate_image(prompt):
     image = pipe(prompt).images[0]