Update app.py

app.py

@@ -1,5 +1,5 @@
-import os
-import spaces
+# import os
+# import spaces
 from dataclasses import dataclass
 
 import gradio as gr
@@ -730,7 +730,7 @@ print(result)
 # model = Flux().to(dtype=torch.bfloat16, device="cuda")
 # result = model.load_state_dict(load_file("/storage/dev/nyanko/flux-dev/flux1-dev.sft"))
 
-@spaces.GPU
+# @spaces.GPU
 @torch.inference_mode()
 def generate_image(
     prompt, width, height, guidance, seed,
@@ -767,12 +767,13 @@ def generate_image(
         timesteps = timesteps[t_idx:]
         x = t * x + (1.0 - t) * init_image.to(x.dtype)
 
-    inp = prepare(t5=t5, clip=clip, img=x, prompt=prompt)
-    x = denoise(model, **inp, timesteps=timesteps, guidance=guidance)
-    x = unpack(x.float(), height, width)
-    with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16):
-        x = x = (x / ae.config.scaling_factor) + ae.config.shift_factor 
-        x = ae.decode(x).sample
+    with torch_device:
+        inp = prepare(t5=t5, clip=clip, img=x, prompt=prompt)
+        x = denoise(model, **inp, timesteps=timesteps, guidance=guidance)
+        x = unpack(x.float(), height, width)
+        with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16):
+            x = x = (x / ae.config.scaling_factor) + ae.config.shift_factor 
+            x = ae.decode(x).sample
 
     x = x.clamp(-1, 1)
     x = rearrange(x[0], "c h w -> h w c")