added model along with final app

.gitmodules

@@ -0,0 +1,5 @@
+[submodule "stylegan3"]
+	path = stylegan3
+	url = https://github.com/NVlabs/stylegan3
+
+    

app.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+from os import pipe
+
+import gradio as gr
+import numpy as np
+from model import Model
+from typing import Tuple
+from diffusers import StableDiffusionPipeline
+import gradio as gr
+from story_generator import StoryGenerator
+import torch
+
+TITLE = ''
+DESCRIPTION = '''# StyleGAN3
+This is an unofficial demo for [https://github.com/NVlabs/stylegan3](https://github.com/NVlabs/stylegan3).
+'''
+
+
+model = Model()
+model_id = "runwayml/stable-diffusion-v1-5"
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+pipe = pipe.to("cuda")  # Remove or comment out this line if using CPU
+
+sg = None
+
+with gr.Blocks(css='style.css') as image_gen_block:
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Tabs():
+        with gr.TabItem('Character'):
+            with gr.Row():
+                with gr.Column():
+                    model_name = gr.Dropdown(list(model.MODEL_NAME_DICT.keys()),
+                                             value='FemaleHero-256-T',
+                                             label='Model')
+                    seed = gr.Slider(0,
+                                     np.iinfo(np.uint32).max,
+                                     step=1,
+                                     value=0,
+                                     label='Seed')
+                    psi = gr.Slider(0,
+                                    2,
+                                    step=0.05,
+                                    value=0.7,
+                                    label='Truncation psi')
+                    tx = gr.Slider(-1,
+                                   1,
+                                   step=0.05,
+                                   value=0,
+                                   label='Translate X')
+                    ty = gr.Slider(-1,
+                                   1,
+                                   step=0.05,
+                                   value=0,
+                                   label='Translate Y')
+                    angle = gr.Slider(-180,
+                                      180,
+                                      step=5,
+                                      value=0,
+                                      label='Angle')
+                    run_button = gr.Button('Run')
+                with gr.Column():
+                    result = gr.Image(label='Result', elem_id='result')
+
+                        # City generation tab
+        with gr.TabItem('City'):
+            with gr.Row():
+                generate_city_button = gr.Button("Generate City")
+            with gr.Row():
+                city_output = gr.Image(label="Generated City", elem_id="city_output")
+
+        with gr.TabItem('Story'):
+            with gr.Row():
+                api_key_input_sg = gr.Textbox(label="OpenAI API Key")
+                prompt_input = gr.Textbox(label='Prompt')
+            generate_button = gr.Button('Generate Story')
+
+            with gr.Row():
+                story_output = gr.Textbox(label='Generated Story',
+                                          placeholder='Click "Generate Story" to see the story',
+                                          readonly=True)
+
+
+
+    def generate_story(prompt, api_key):
+        sg = StoryGenerator(api_key)
+        return sg.generate_story(prompt)
+
+    def generate_city():
+        model_id = "runwayml/stable-diffusion-v1-5"
+        pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+        pipe = pipe.to("cuda")
+        city_prompt = f"A metropolis city, HD"
+        city_image = pipe(city_prompt).images[0]
+        return city_image
+
+    generate_city_button.click(fn=generate_city, inputs=None, outputs=city_output)
+    model_name.change(fn=model.set_model, inputs=model_name, outputs=None)
+    run_button.click(fn=model.generate_image,
+                    inputs=[
+                        model_name,
+                        seed,
+                        psi,
+                        tx,
+                        ty,
+                        angle,
+                    ],
+                    outputs=result)
+    generate_button.click(fn=generate_story, inputs=[prompt_input, api_key_input_sg], outputs=story_output)
+
+
+
+
+image_gen_block.queue().launch(show_api=False)
+
+

model.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import os
+import pathlib
+import pickle
+import sys
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+current_dir = pathlib.Path(__file__).parent
+submodule_dir = current_dir / 'stylegan3'
+sys.path.insert(0, submodule_dir.as_posix())
+
+class Model:
+    MODEL_NAME_DICT = {'FemaleHero-256-T': 'network-snapshot-000120.pkl'}
+
+    def __init__(self):
+        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        self._download_all_models()
+        self.model_name = 'FemaleHero-256-T'
+        self.model = self._load_model(self.model_name)
+
+    def _load_model(self, model_name: str) -> nn.Module:
+        file_name = self.MODEL_NAME_DICT[model_name]
+        with open(file_name, 'rb') as f:
+            model = pickle.load(f)['G_ema']
+        model.eval()
+        model.to(self.device)
+        return model
+
+    def set_model(self, model_name: str) -> None:
+        if model_name == self.model_name:
+            return
+        self.model_name = model_name
+        self.model = self._load_model(model_name)
+
+    def _download_all_models(self):
+        pass
+
+    @staticmethod
+    def make_transform(translate: tuple[float, float], angle: float) -> np.ndarray:
+        mat = np.eye(3)
+        sin = np.sin(angle / 360 * np.pi * 2)
+        cos = np.cos(angle / 360 * np.pi * 2)
+        mat[0][0] = cos
+        mat[0][1] = sin
+        mat[0][2] = translate[0]
+        mat[1][0] = -sin
+        mat[1][1] = cos
+        mat[1][2] = translate[1]
+        return mat
+
+    def generate_z(self, seed: int) -> torch.Tensor:
+        seed = int(np.clip(seed, 0, np.iinfo(np.uint32).max))
+        z = np.random.RandomState(seed).randn(1, self.model.z_dim)
+        return torch.from_numpy(z).float().to(self.device)
+
+    def postprocess(self, tensor: torch.Tensor) -> np.ndarray:
+        tensor = (tensor.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+        return tensor.cpu().numpy()
+
+    def set_transform(self, tx: float, ty: float, angle: float) -> None:
+        mat = self.make_transform((tx, ty), angle)
+        mat = np.linalg.inv(mat)
+        self.model.synthesis.input.transform.copy_(torch.from_numpy(mat))
+
+    @torch.inference_mode()
+    def generate(self, z: torch.Tensor, label: torch.Tensor, truncation_psi: float) -> torch.Tensor:
+        return self.model(z, label, truncation_psi=truncation_psi)
+
+    def generate_image(self, model_name: str, seed: int, truncation_psi: float, tx: float, ty: float, angle: float) -> np.ndarray:
+        self.set_model(model_name)
+        self.set_transform(tx, ty, angle)
+        z = self.generate_z(seed)
+        label = torch.zeros([1, self.model.c_dim], device=self.device)
+        out = self.generate(z, label, truncation_psi=truncation_psi)
+        out = self.postprocess(out)
+        return out[0]

network-snapshot-000120.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:494b89ca839757d1c76a9a1b5eb0be5ce01e0d66eee8a55e1563deba9734c7a6
+size 343475301

story_generator.py

@@ -0,0 +1,20 @@
+import openai
+
+class StoryGenerator:
+    def __init__(self, api_key):
+
+        openai.api_key = api_key
+        self.completion_model = "text-davinci-002"
+
+
+    def generate_story(self, prompt):
+        prompt = f"Once upon a time, {prompt}"
+        response = openai.Completion.create(
+            engine=self.completion_model,
+            prompt=prompt,
+            max_tokens=1024,
+            n=1,
+            stop=None,
+            temperature=0.5,
+        )
+        return response.choices[0].text

style.css

@@ -0,0 +1,12 @@
+h1 {
+    text-align: center;
+  }
+  div#result {
+    max-width: 600px;
+    max-height: 600px;
+  }
+  img#visitor-badge {
+    display: block;
+    margin: auto;
+  }
+