Update theme.py

theme.py

@@ -1,67 +1,68 @@
-import random
-import replicate
-import base64
-import os
-
-os.environ["REPLICATE_API_TOKEN"] = os.getenv("REPLICATE_API_TOKEN")
-
-
-def create_flux_request(prompt_for_image_generation):
-    payload = {
-        "prompt": prompt_for_image_generation,
-        "guidance": 3.5,
-        "num_outputs": 1,
-        "aspect_ratio": "3:4"
-    }
-    output = replicate.run(
-        "black-forest-labs/flux-dev",
-        input=payload
-    )
-    return output
-
-
-def flux_generated_image(prompt_for_image_generation):
-    try:
-        flux_response_object = create_flux_request(prompt_for_image_generation)
-        data_uri = flux_response_object[0].url
-        header, encoded = data_uri.split(',', 1)
-        file_data = base64.b64decode(encoded)
-        random_int_for_file_prefix = random.randint(1, 1000000)
-        output_image_file_name = f"{random_int_for_file_prefix}_ide_theme_image.png"
-        with open(output_image_file_name, "wb") as f:
-            f.write(file_data)
-        return {"success": True, "file_name": output_image_file_name}
-    except Exception as e:
-        return {"success": False, "error": e}
-
-
-def create_flux_request_seed(prompt_for_image_generation, seed, aspect_ratio):
-    print(f"YE SEED HAI MEPE:{seed}")
-    payload = {
-        "prompt": prompt_for_image_generation,
-        "guidance": 3.5,
-        "num_outputs": 1,
-        "aspect_ratio": str(aspect_ratio),
-        "seed": int(seed)
-
-    }
-    output = replicate.run(
-        "black-forest-labs/flux-dev",
-        input=payload
-    )
-    return output
-
-
-def flux_generated_image_seed(prompt_for_image_generation, seed, aspect_ratio):
-    try:
-        flux_response_object = create_flux_request_seed(prompt_for_image_generation, seed, aspect_ratio)
-        data_uri = flux_response_object[0].url
-        header, encoded = data_uri.split(',', 1)
-        file_data = base64.b64decode(encoded)
-        random_int_for_file_prefix = random.randint(1, 1000000)
-        output_image_file_name = f"{random_int_for_file_prefix}_ide_theme_image.png"
-        with open(output_image_file_name, "wb") as f:
-            f.write(file_data)
-        return {"success": True, "file_name": output_image_file_name}
-    except Exception as e:
-        return {"success": False, "error": e}
+import random
+import replicate
+import base64
+import os
+
+os.environ["REPLICATE_API_TOKEN"] = os.getenv("REPLICATE_API_TOKEN")
+
+
+def create_flux_request(prompt_for_image_generation):
+    payload = {
+        "prompt": prompt_for_image_generation,
+        "guidance": 3.5,
+        "num_outputs": 1,
+        "aspect_ratio": "3:4"
+    }
+    output = replicate.run(
+        "black-forest-labs/flux-dev",
+        input=payload
+    )
+    return output
+
+
+def flux_generated_image(prompt_for_image_generation):
+    try:
+        p = "Create Non NSFW image." + prompt_for_image_generation
+        flux_response_object = create_flux_request(p)
+        data_uri = flux_response_object[0].url
+        header, encoded = data_uri.split(',', 1)
+        file_data = base64.b64decode(encoded)
+        random_int_for_file_prefix = random.randint(1, 1000000)
+        output_image_file_name = f"{random_int_for_file_prefix}_ide_theme_image.png"
+        with open(output_image_file_name, "wb") as f:
+            f.write(file_data)
+        return {"success": True, "file_name": output_image_file_name}
+    except Exception as e:
+        return {"success": False, "error": e}
+
+
+def create_flux_request_seed(prompt_for_image_generation, seed, aspect_ratio):
+    print(f"YE SEED HAI MEPE:{seed}")
+    payload = {
+        "prompt": prompt_for_image_generation,
+        "guidance": 3.5,
+        "num_outputs": 1,
+        "aspect_ratio": str(aspect_ratio),
+        "seed": int(seed)
+
+    }
+    output = replicate.run(
+        "black-forest-labs/flux-dev",
+        input=payload
+    )
+    return output
+
+
+def flux_generated_image_seed(prompt_for_image_generation, seed, aspect_ratio):
+    try:
+        flux_response_object = create_flux_request_seed(prompt_for_image_generation, seed, aspect_ratio)
+        data_uri = flux_response_object[0].url
+        header, encoded = data_uri.split(',', 1)
+        file_data = base64.b64decode(encoded)
+        random_int_for_file_prefix = random.randint(1, 1000000)
+        output_image_file_name = f"{random_int_for_file_prefix}_ide_theme_image.png"
+        with open(output_image_file_name, "wb") as f:
+            f.write(file_data)
+        return {"success": True, "file_name": output_image_file_name}
+    except Exception as e:
+        return {"success": False, "error": e}