Upload 2 files

utils.py

@@ -7,7 +7,6 @@ import torch
 import uuid
 from PIL import Image, PngImagePlugin
 from datetime import datetime
-from dataclasses import dataclass
 from typing import Callable, Dict, Optional, Tuple, Any, List
 from diffusers import (
     DDIMScheduler,
@@ -22,18 +21,8 @@ import logging
 
 MAX_SEED = np.iinfo(np.int32).max
 
-
-@dataclass
-class StyleConfig:
-    prompt: str
-    negative_prompt: str
-
-
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
-
+def is_space_environment():
+    return "SPACE_ID" in os.environ and os.environ.get("SYSTEM") == "spaces"
 
 def seed_everything(seed: int) -> torch.Generator:
     torch.manual_seed(seed)
@@ -217,10 +206,19 @@ def load_pipeline(model_name: str, device: torch.device, hf_token: Optional[str]
                 add_watermarker=False
             )
 
-        pipe.to(device)
+        # デバイス移動の部分を修正
+        if "SPACE_ID" in os.environ and os.environ.get("SYSTEM") == "spaces":
+            # Stateless GPU環境ではデバイス移動を特別に扱う
+            return pipe
+        else:
+            # 通常の環境では以前のコードを使用
+            pipe.to(device)
+            return pipe
+
         logging.info("Pipeline loaded successfully!")
         return pipe
     except Exception as e:
         logging.error(f"Failed to load pipeline: {str(e)}", exc_info=True)
         raise
 
+